fix: track transaction number in changes (#33)

matter-labs · May 23, 2024 · e683ae8 · e683ae8
1 parent d9cb911
commit e683ae8
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/src/instruction_handlers/precompiles.rs b/src/instruction_handlers/precompiles.rs
@@ -50,6 +50,7 @@ fn precompile_call(
         let query = LogQuery {
             timestamp: Timestamp(0),
             key: abi.to_u256(),
+            // only two first fields are read by the precompile
             tx_number_in_block: Default::default(),
             aux_byte: Default::default(),
             shard_id: Default::default(),

diff --git a/src/instruction_handlers/storage.rs b/src/instruction_handlers/storage.rs
@@ -27,9 +27,13 @@ fn sstore(
             let key = Register1::get(args, &mut vm.state);
             let value = Register2::get(args, &mut vm.state);
 
-            let refund =
-                vm.world_diff
-                    .write_storage(world, vm.state.current_frame.address, key, value);
+            let refund = vm.world_diff.write_storage(
+                world,
+                vm.state.current_frame.address,
+                key,
+                value,
+                vm.state.transaction_number,
+            );
 
             assert!(refund <= SSTORE_COST);
             vm.state.current_frame.gas += refund;

diff --git a/src/world_diff.rs b/src/world_diff.rs
@@ -15,7 +15,7 @@ use zkevm_opcode_defs::system_params::{
 #[derive(Default)]
 pub struct WorldDiff {
     // These are rolled back on revert or panic (and when the whole VM is rolled back).
-    storage_changes: RollbackableMap<(H160, U256), U256>,
+    storage_changes: RollbackableMap<(H160, U256), (u16, U256)>,
     paid_changes: RollbackableMap<(H160, U256), u32>,
     transient_storage_changes: RollbackableMap<(H160, U256), U256>,
     events: RollbackableLog<Event>,
@@ -68,6 +68,7 @@ impl WorldDiff {
             .as_ref()
             .get(&(contract, key))
             .cloned()
+            .map(|v| v.1)
             .unwrap_or_else(|| world.read_storage(contract, key));
 
         let refund = if world.is_free_storage_slot(&contract, &key)
@@ -89,8 +90,10 @@ impl WorldDiff {
         contract: H160,
         key: U256,
         value: U256,
+        tx_number: u16,
     ) -> u32 {
-        self.storage_changes.insert((contract, key), value);
+        self.storage_changes
+            .insert((contract, key), (tx_number, value));
 
         if world.is_free_storage_slot(&contract, &key) {
             return WARM_WRITE_REFUND;
@@ -124,18 +127,22 @@ impl WorldDiff {
         refund
     }
 
-    pub fn get_storage_state(&self) -> &BTreeMap<(H160, U256), U256> {
+    pub fn get_storage_state(&self) -> &BTreeMap<(H160, U256), (u16, U256)> {
         self.storage_changes.as_ref()
     }
 
-    pub fn get_storage_changes(&self) -> BTreeMap<(H160, U256), (Option<U256>, U256)> {
+    #[allow(clippy::type_complexity)]
+    pub fn get_storage_changes(
+        &self,
+    ) -> BTreeMap<(H160, U256), (Option<(u16, U256)>, (u16, U256))> {
         self.storage_changes.changes_after(0)
     }
 
+    #[allow(clippy::type_complexity)]
     pub fn get_storage_changes_after(
         &self,
         snapshot: &Snapshot,
-    ) -> BTreeMap<(H160, U256), (Option<U256>, U256)> {
+    ) -> BTreeMap<(H160, U256), (Option<(u16, U256)>, (u16, U256))> {
         self.storage_changes.changes_after(snapshot.storage_changes)
     }