feat: Track storage_refunds and pubdata_costs stats (#48)

# What ❔

- Tracks `storage_refunds` and `pubdata_costs` in `WorldDiff` in order
to fill in the corresponding data in the VM glue code.
- Fixes decommit opcode semantics, refunding extra gas cost if the
decommitted bytecode is not fresh (i.e., was decommitted previously).

## Why ❔

- `storage_refunds` and `pubdata_costs` data is not read directly on the
EN, but is persisted for each batch in the state keeper. By design, this
data is expected to filled in by the state keeper, so filling it
retrospectively could be a problem.
- Divergency in decommit opcode semantics is easily observable (e.g., by
gas costs).

fix: Fix decommit opcode semantics

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