Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rerragange sov-modules-api scratchpad for better readability #953

Merged
merged 1 commit into from
Sep 28, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
221 changes: 108 additions & 113 deletions module-system/sov-modules-api/src/state/scratchpad.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,53 @@ impl<S: Storage> StateReaderAndWriter for Delta<S> {

type RevertableWrites = HashMap<CacheKey, Option<CacheValue>>;

struct AccessoryDelta<S: Storage> {
// This inner storage is never accessed inside the zkVM because reads are
// not allowed, so it can result as dead code.
#[allow(dead_code)]
storage: S,
writes: RevertableWrites,
}

impl<S: Storage> AccessoryDelta<S> {
fn new(storage: S) -> Self {
Self {
storage,
writes: Default::default(),
}
}

fn freeze(&mut self) -> OrderedReadsAndWrites {
let mut reads_and_writes = OrderedReadsAndWrites::default();
let writes = std::mem::take(&mut self.writes);

for write in writes {
reads_and_writes.ordered_writes.push((write.0, write.1));
}

reads_and_writes
}
}

impl<S: Storage> StateReaderAndWriter for AccessoryDelta<S> {
fn get(&mut self, key: &StorageKey) -> Option<StorageValue> {
let cache_key = key.to_cache_key();
if let Some(value) = self.writes.get(&cache_key) {
return value.clone().map(Into::into);
}
self.storage.get_accessory(key)
}

fn set(&mut self, key: &StorageKey, value: StorageValue) {
self.writes
.insert(key.to_cache_key(), Some(value.into_cache_value()));
}

fn delete(&mut self, key: &StorageKey) {
self.writes.insert(key.to_cache_key(), None);
}
}

/// This structure is responsible for storing the `read-write` set.
///
/// A [`StateCheckpoint`] can be obtained from a [`WorkingSet`] in two ways:
Expand All @@ -80,11 +127,6 @@ impl<C: Context> StateCheckpoint<C> {
}
}

/// Fetches a value from the underlying storage.
pub fn get(&mut self, key: &StorageKey) -> Option<StorageValue> {
self.delta.get(key)
}

/// Creates a new [`StateCheckpoint`] instance without any changes, backed
/// by the given [`Storage`] and witness.
pub fn with_witness(
Expand Down Expand Up @@ -131,61 +173,82 @@ impl<C: Context> StateCheckpoint<C> {
}
}

struct AccessoryDelta<S: Storage> {
// This inner storage is never accessed inside the zkVM because reads are
// not allowed, so it can result as dead code.
#[allow(dead_code)]
storage: S,
writes: RevertableWrites,
/// This structure contains the read-write set and the events collected during the execution of a transaction.
/// There are two ways to convert it into a StateCheckpoint:
/// 1. By using the checkpoint() method, where all the changes are added to the underlying StateCheckpoint.
/// 2. By using the revert method, where the most recent changes are reverted and the previous `StateCheckpoint` is returned.
pub struct WorkingSet<C: Context> {
delta: RevertableWriter<Delta<C::Storage>>,
accessory_delta: RevertableWriter<AccessoryDelta<C::Storage>>,
events: Vec<Event>,
}

impl<S: Storage> AccessoryDelta<S> {
fn new(storage: S) -> Self {
Self {
storage,
writes: Default::default(),
}
impl<C: Context> WorkingSet<C> {
/// Creates a new [`WorkingSet`] instance backed by the given [`Storage`].
///
/// The witness value is set to [`Default::default`]. Use
/// [`WorkingSet::with_witness`] to set a custom witness value.
pub fn new(inner: <C as Spec>::Storage) -> Self {
StateCheckpoint::new(inner).to_revertable()
}

fn freeze(&mut self) -> OrderedReadsAndWrites {
let mut reads_and_writes = OrderedReadsAndWrites::default();
let writes = std::mem::take(&mut self.writes);
/// Returns a handler for the accessory state (non-JMT state).
///
/// You can use this method when calling getters and setters on accessory
/// state containers, like [`AccessoryStateMap`](crate::AccessoryStateMap).
pub fn accessory_state(&mut self) -> AccessoryWorkingSet<C> {
AccessoryWorkingSet { ws: self }
}

for write in writes {
reads_and_writes.ordered_writes.push((write.0, write.1));
}
/// Creates a new [`WorkingSet`] instance backed by the given [`Storage`]
/// and a custom witness value.
pub fn with_witness(
inner: <C as Spec>::Storage,
witness: <<C as Spec>::Storage as Storage>::Witness,
) -> Self {
StateCheckpoint::with_witness(inner, witness).to_revertable()
}

reads_and_writes
/// Turns this [`WorkingSet`] into a [`StateCheckpoint`], in preparation for
/// committing the changes to the underlying [`Storage`] via
/// [`StateCheckpoint::freeze`].
pub fn checkpoint(self) -> StateCheckpoint<C> {
StateCheckpoint {
delta: self.delta.commit(),
accessory_delta: self.accessory_delta.commit(),
}
}
}

impl<S: Storage> StateReaderAndWriter for AccessoryDelta<S> {
fn get(&mut self, key: &StorageKey) -> Option<StorageValue> {
let cache_key = key.to_cache_key();
if let Some(value) = self.writes.get(&cache_key) {
return value.clone().map(Into::into);
/// Reverts the most recent changes to this [`WorkingSet`], returning a pristine
/// [`StateCheckpoint`] instance.
pub fn revert(self) -> StateCheckpoint<C> {
StateCheckpoint {
delta: self.delta.revert(),
accessory_delta: self.accessory_delta.revert(),
}
self.storage.get_accessory(key)
}

fn set(&mut self, key: &StorageKey, value: StorageValue) {
self.writes
.insert(key.to_cache_key(), Some(value.into_cache_value()));
/// Adds an event to the working set.
pub fn add_event(&mut self, key: &str, value: &str) {
self.events.push(Event::new(key, value));
}

fn delete(&mut self, key: &StorageKey) {
self.writes.insert(key.to_cache_key(), None);
/// Extracts all events from this working set.
pub fn take_events(&mut self) -> Vec<Event> {
std::mem::take(&mut self.events)
}
}

/// This structure contains the read-write set and the events collected during the execution of a transaction.
/// There are two ways to convert it into a StateCheckpoint:
/// 1. By using the checkpoint() method, where all the changes are added to the underlying StateCheckpoint.
/// 2. By using the revert method, where the most recent changes are reverted and the previous `StateCheckpoint` is returned.
pub struct WorkingSet<C: Context> {
delta: RevertableWriter<Delta<C::Storage>>,
accessory_delta: RevertableWriter<AccessoryDelta<C::Storage>>,
events: Vec<Event>,
/// Returns an immutable slice of all events that have been previously
/// written to this working set.
pub fn events(&self) -> &[Event] {
&self.events
}

/// Returns an immutable reference to the [`Storage`] instance backing this
/// working set.
pub fn backing(&self) -> &<C as Spec>::Storage {
&self.delta.inner.inner
}
}

impl<C: Context> StateReaderAndWriter for WorkingSet<C> {
Expand Down Expand Up @@ -286,74 +349,6 @@ impl<T: StateReaderAndWriter> StateReaderAndWriter for RevertableWriter<T> {
}
}

impl<C: Context> WorkingSet<C> {
/// Creates a new [`WorkingSet`] instance backed by the given [`Storage`].
///
/// The witness value is set to [`Default::default`]. Use
/// [`WorkingSet::with_witness`] to set a custom witness value.
pub fn new(inner: <C as Spec>::Storage) -> Self {
StateCheckpoint::new(inner).to_revertable()
}

/// Returns a handler for the accessory state (non-JMT state).
///
/// You can use this method when calling getters and setters on accessory
/// state containers, like [`AccessoryStateMap`](crate::AccessoryStateMap).
pub fn accessory_state(&mut self) -> AccessoryWorkingSet<C> {
AccessoryWorkingSet { ws: self }
}

/// Creates a new [`WorkingSet`] instance backed by the given [`Storage`]
/// and a custom witness value.
pub fn with_witness(
inner: <C as Spec>::Storage,
witness: <<C as Spec>::Storage as Storage>::Witness,
) -> Self {
StateCheckpoint::with_witness(inner, witness).to_revertable()
}

/// Turns this [`WorkingSet`] into a [`StateCheckpoint`], in preparation for
/// committing the changes to the underlying [`Storage`] via
/// [`StateCheckpoint::freeze`].
pub fn checkpoint(self) -> StateCheckpoint<C> {
StateCheckpoint {
delta: self.delta.commit(),
accessory_delta: self.accessory_delta.commit(),
}
}

/// Reverts the most recent changes to this [`WorkingSet`], returning a pristine
/// [`StateCheckpoint`] instance.
pub fn revert(self) -> StateCheckpoint<C> {
StateCheckpoint {
delta: self.delta.revert(),
accessory_delta: self.accessory_delta.revert(),
}
}

/// Adds an event to the working set.
pub fn add_event(&mut self, key: &str, value: &str) {
self.events.push(Event::new(key, value));
}

/// Extracts all events from this working set.
pub fn take_events(&mut self) -> Vec<Event> {
std::mem::take(&mut self.events)
}

/// Returns an immutable slice of all events that have been previously
/// written to this working set.
pub fn events(&self) -> &[Event] {
&self.events
}

/// Returns an immutable reference to the [`Storage`] instance backing this
/// working set.
pub fn backing(&self) -> &<C as Spec>::Storage {
&self.delta.inner.inner
}
}

pub(crate) trait StateReaderAndWriter {
fn get(&mut self, key: &StorageKey) -> Option<StorageValue>;

Expand Down