Skip to content

Commit

Permalink
Delta's go first, remove get from StateCheckpoint (#953)
Browse files Browse the repository at this point in the history
  • Loading branch information
@citizen-stig @dubbelosix
citizen-stig authored and dubbelosix committed Sep 29, 2023
1 parent 1e0c87c commit 1d8ac46
Showing 1 changed file with 108 additions and 113 deletions.
221 changes: 108 additions & 113 deletions module-system/sov-modules-api/src/state/scratchpad.rs
View file
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

0 comments on commit 1d8ac46

Please sign in to comment.