add method for getting storage changes after a snapshot

matter-labs · May 8, 2024 · b9c2e44 · b9c2e44
1 parent b3e07ec
commit b9c2e44
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/src/modified_world.rs b/src/modified_world.rs
@@ -127,6 +127,13 @@ impl ModifiedWorld {
         self.storage_changes.as_ref()
     }
 
+    pub fn get_storage_changes_after(
+        &self,
+        snapshot: &Snapshot,
+    ) -> BTreeMap<(H160, U256), (Option<U256>, U256)> {
+        self.storage_changes.changes_after(snapshot.storage_changes)
+    }
+
     pub(crate) fn record_event(&mut self, event: Event) {
         self.events.push(event);
     }

diff --git a/src/rollback.rs b/src/rollback.rs
@@ -20,6 +20,20 @@ impl<K: Ord + Clone, V: Clone> RollbackableMap<K, V> {
         self.old_entries.push((key.clone(), old_value.clone()));
         old_value
     }
+
+    pub(crate) fn changes_after(
+        &self,
+        snapshot: <Self as Rollback>::Snapshot,
+    ) -> BTreeMap<K, (Option<V>, V)> {
+        let mut changes = BTreeMap::new();
+        for (key, old_value) in self.old_entries[snapshot..].iter().rev() {
+            changes
+                .entry(key.clone())
+                .and_modify(|(old, _)| *old = old_value.clone())
+                .or_insert((old_value.clone(), self.map.get(key).unwrap().clone()));
+        }
+        changes
+    }
 }
 
 impl<K: Ord, V> Rollback for RollbackableMap<K, V> {