feat: Track storage_refunds and pubdata_costs stats

src/decommit.rs

@@ -19,8 +19,8 @@ impl WorldDiff {
         let mut is_evm = false;
 
         let mut code_info = {
-            let (code_info, _) =
-                self.read_storage(world, deployer_system_contract_address, address);
+            let code_info =
+                self.read_storage_without_refund(world, deployer_system_contract_address, address);
             let mut code_info_bytes = [0; 32];
             code_info.to_big_endian(&mut code_info_bytes);
 
@@ -77,11 +77,12 @@ impl WorldDiff {
         Some((UnpaidDecommit { cost, code_key }, is_evm))
     }
 
-    pub(crate) fn decommit_opcode(&mut self, world: &mut dyn World, code_hash: U256) -> Vec<u8> {
-        let mut code_info_bytes = [0; 32];
-        code_hash.to_big_endian(&mut code_info_bytes);
-        self.decommitted_hashes.insert(code_hash, ());
-        world.decommit_code(code_hash)
+    /// Returns the decommitted contract code and a flag set to `true` if this is a fresh decommit (i.e.,
+    /// the code wasn't decommitted previously in the same VM run).
+    #[doc(hidden)] // should be used for testing purposes only; can break VM operation otherwise
+    pub fn decommit_opcode(&mut self, world: &mut dyn World, code_hash: U256) -> (Vec<u8>, bool) {
+        let was_decommitted = self.decommitted_hashes.insert(code_hash, ()).is_some();
+        (world.decommit_code(code_hash), !was_decommitted)
     }
 
     pub(crate) fn pay_for_decommit(

src/instruction_handlers/decommit.rs

@@ -33,7 +33,10 @@ fn decommit(
             return;
         }
 
-        let program = vm.world_diff.decommit_opcode(world, code_hash);
+        let (program, is_fresh) = vm.world_diff.decommit_opcode(world, code_hash);
+        if !is_fresh {
+            vm.state.current_frame.gas += extra_cost;
+        }
 
         let heap = vm.state.heaps.allocate();
         vm.state.current_frame.heaps_i_am_keeping_alive.push(heap);
@@ -49,6 +52,7 @@ fn decommit(
         Register1::set_fat_ptr(args, &mut vm.state, value);
     })
 }
+
 impl Instruction {
     pub fn from_decommit(
         abi: Register1,

src/world_diff.rs

@@ -20,7 +20,9 @@ pub struct WorldDiff {
     transient_storage_changes: RollbackableMap<(H160, U256), U256>,
     events: RollbackableLog<Event>,
     l2_to_l1_logs: RollbackableLog<L2ToL1Log>,
-    pub pubdata: RollbackablePod<i32>,
+    pub(crate) pubdata: RollbackablePod<i32>,
+    storage_refunds: RollbackableLog<u32>,
+    pubdata_costs: RollbackableLog<i32>,
 
     // The fields below are only rolled back when the whole VM is rolled back.
     pub(crate) decommitted_hashes: RollbackableSet<U256>,
@@ -36,6 +38,8 @@ pub struct ExternalSnapshot {
     pub(crate) decommitted_hashes: <RollbackableMap<U256, ()> as Rollback>::Snapshot,
     read_storage_slots: <RollbackableMap<(H160, U256), ()> as Rollback>::Snapshot,
     written_storage_slots: <RollbackableMap<(H160, U256), ()> as Rollback>::Snapshot,
+    storage_refunds: <RollbackableLog<u32> as Rollback>::Snapshot,
+    pubdata_costs: <RollbackableLog<i32> as Rollback>::Snapshot,
 }
 
 /// There is no address field because nobody is interested in events that don't come
@@ -66,12 +70,35 @@ impl WorldDiff {
         world: &mut dyn World,
         contract: H160,
         key: U256,
+    ) -> (U256, u32) {
+        let (value, refund) = self.read_storage_inner(world, contract, key);
+        self.storage_refunds.push(refund);
+        (value, refund)
+    }
+
+    /// Same as [`Self::read_storage()`], but without recording the refund value (which is important
+    /// because the storage is read not only from the `sload` op handler, but also from the `farcall` op handler;
+    /// the latter must not record a refund as per previous VM versions).
+    pub(crate) fn read_storage_without_refund(
+        &mut self,
+        world: &mut dyn World,
+        contract: H160,
+        key: U256,
+    ) -> U256 {
+        self.read_storage_inner(world, contract, key).0
+    }
+
+    fn read_storage_inner(
+        &mut self,
+        world: &mut dyn World,
+        contract: H160,
+        key: U256,
     ) -> (U256, u32) {
         let value = self
             .storage_changes
             .as_ref()
             .get(&(contract, key))
-            .cloned()
+            .copied()
             .unwrap_or_else(|| world.read_storage(contract, key).unwrap_or_default());
 
         let refund = if world.is_free_storage_slot(&contract, &key)
@@ -82,7 +109,7 @@ impl WorldDiff {
             self.read_storage_slots.add((contract, key));
             0
         };
-
+        self.pubdata_costs.push(0);
         (value, refund)
     }
 
@@ -102,6 +129,8 @@ impl WorldDiff {
             .or_insert_with(|| world.read_storage(contract, key));
 
         if world.is_free_storage_slot(&contract, &key) {
+            self.storage_refunds.push(WARM_WRITE_REFUND);
+            self.pubdata_costs.push(0);
             return WARM_WRITE_REFUND;
         }
 
@@ -128,11 +157,25 @@ impl WorldDiff {
             }
         };
 
-        self.pubdata.0 += (update_cost as i32) - (prepaid as i32);
-
+        let pubdata_cost = (update_cost as i32) - (prepaid as i32);
+        self.pubdata.0 += pubdata_cost;
+        self.storage_refunds.push(refund);
+        self.pubdata_costs.push(pubdata_cost);
         refund
     }
 
+    pub fn pubdata(&self) -> i32 {
+        self.pubdata.0
+    }
+
+    pub fn storage_refunds(&self) -> &[u32] {
+        self.storage_refunds.as_ref()
+    }
+
+    pub fn pubdata_costs(&self) -> &[i32] {
+        self.pubdata_costs.as_ref()
+    }
+
     pub fn get_storage_state(&self) -> &BTreeMap<(H160, U256), U256> {
         self.storage_changes.as_ref()
     }
@@ -173,17 +216,16 @@ impl WorldDiff {
     }
 
     pub(crate) fn read_transient_storage(&mut self, contract: H160, key: U256) -> U256 {
-        let value = self
-            .transient_storage_changes
+        self.pubdata_costs.push(0);
+        self.transient_storage_changes
             .as_ref()
             .get(&(contract, key))
-            .cloned()
-            .unwrap_or_default();
-
-        value
+            .copied()
+            .unwrap_or_default()
     }
 
     pub(crate) fn write_transient_storage(&mut self, contract: H160, key: U256, value: U256) {
+        self.pubdata_costs.push(0);
         self.transient_storage_changes
             .insert((contract, key), value);
     }
@@ -212,8 +254,9 @@ impl WorldDiff {
         self.l2_to_l1_logs.logs_after(snapshot.l2_to_l1_logs)
     }
 
-    pub fn get_decommitted_hashes(&self) -> &BTreeMap<U256, ()> {
-        self.decommitted_hashes.as_ref()
+    /// Returns hashes of decommitted contract bytecodes in no particular order.
+    pub fn decommitted_hashes(&self) -> impl Iterator<Item = U256> + '_ {
+        self.decommitted_hashes.as_ref().keys().copied()
     }
 
     /// Get a snapshot for selecting which logs [Self::events_after] & Co output.
@@ -230,24 +273,14 @@ impl WorldDiff {
         }
     }
 
-    pub(crate) fn rollback(
-        &mut self,
-        Snapshot {
-            storage_changes,
-            paid_changes,
-            events,
-            l2_to_l1_logs,
-            transient_storage_changes,
-            pubdata,
-        }: Snapshot,
-    ) {
-        self.storage_changes.rollback(storage_changes);
-        self.paid_changes.rollback(paid_changes);
-        self.events.rollback(events);
-        self.l2_to_l1_logs.rollback(l2_to_l1_logs);
+    pub(crate) fn rollback(&mut self, snapshot: Snapshot) {
+        self.storage_changes.rollback(snapshot.storage_changes);
+        self.paid_changes.rollback(snapshot.paid_changes);
+        self.events.rollback(snapshot.events);
+        self.l2_to_l1_logs.rollback(snapshot.l2_to_l1_logs);
         self.transient_storage_changes
-            .rollback(transient_storage_changes);
-        self.pubdata.rollback(pubdata);
+            .rollback(snapshot.transient_storage_changes);
+        self.pubdata.rollback(snapshot.pubdata);
     }
 
     /// This function must only be called during the initial frame
@@ -265,11 +298,15 @@ impl WorldDiff {
             decommitted_hashes: self.decommitted_hashes.snapshot(),
             read_storage_slots: self.read_storage_slots.snapshot(),
             written_storage_slots: self.written_storage_slots.snapshot(),
+            storage_refunds: self.storage_refunds.snapshot(),
+            pubdata_costs: self.pubdata_costs.snapshot(),
         }
     }
 
     pub(crate) fn external_rollback(&mut self, snapshot: ExternalSnapshot) {
         self.rollback(snapshot.internal_snapshot);
+        self.storage_refunds.rollback(snapshot.storage_refunds);
+        self.pubdata_costs.rollback(snapshot.pubdata_costs);
         self.decommitted_hashes
             .rollback(snapshot.decommitted_hashes);
         self.read_storage_slots