add method for getting storage changes after a snapshot

src/modified_world.rs

@@ -109,6 +109,17 @@ impl ModifiedWorld {
         }
     }
 
+    pub fn get_storage_changes(&self) -> &BTreeMap<(H160, U256), U256> {
+        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 fn prepaid_for_write(&self, address: H160, key: U256) -> u32 {
         self.paid_changes
             .as_ref()
@@ -121,10 +132,6 @@ impl ModifiedWorld {
         self.paid_changes.insert((address, key), price)
     }
 
-    pub fn get_storage_changes(&self) -> &BTreeMap<(H160, U256), U256> {
-        self.storage_changes.as_ref()
-    }
-
     pub(crate) fn record_event(&mut self, event: Event) {
         self.events.push(event);
     }

src/rollback.rs

@@ -14,11 +14,25 @@ pub struct RollbackableMap<K: Ord, V> {
     old_entries: Vec<(K, Option<V>)>,
 }
 
-impl<K: Ord + Clone, V> RollbackableMap<K, V> {
+impl<K: Ord + Clone, V: Clone> RollbackableMap<K, V> {
     pub fn insert(&mut self, key: K, value: V) {
         self.old_entries
             .push((key.clone(), self.map.insert(key, 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> {