tally storage cycles at the end; fix bugs in read_storage_slots

src/decommit.rs

@@ -85,7 +85,7 @@ impl WorldDiff {
         world: &mut impl World<T>,
         code_hash: U256,
     ) -> (Vec<u8>, bool) {
-        let is_new = self.decommitted_hashes.add(code_hash);
+        let is_new = !self.decommitted_hashes.add(code_hash);
         (world.decommit_code(code_hash), is_new)
     }
 
@@ -97,7 +97,7 @@ impl WorldDiff {
     ) -> Option<(Program<T, W>, usize)> {
         // We intentionally record a decommitment event even if actual decommitment never happens because of an out-of-gas error.
         // This is how the old VM behaves.
-        let is_new = self.decommitted_hashes.add(decommit.code_key);
+        let is_new = !self.decommitted_hashes.add(decommit.code_key);
 
         if decommit.cost > *gas {
             // Unlike all other gas costs, this one is not paid if low on gas.

src/instruction_handlers/decommit.rs

@@ -3,7 +3,6 @@ use crate::{
     addressing_modes::{Arguments, Destination, Register1, Register2, Source},
     fat_pointer::FatPointer,
     instruction::ExecutionStatus,
-    vm::STORAGE_READ_STORAGE_APPLICATION_CYCLES,
     Instruction, VirtualMachine, World,
 };
 use eravm_stable_interface::{opcodes, Tracer};
@@ -36,9 +35,6 @@ fn decommit<T: Tracer, W: World<T>>(
         let (program, is_fresh) = vm.world_diff.decommit_opcode(world, code_hash);
         if !is_fresh {
             vm.state.current_frame.gas += extra_cost;
-        } else {
-            vm.state.cycle_counts.storage_application_cycles +=
-                STORAGE_READ_STORAGE_APPLICATION_CYCLES as usize;
         }
 
         let heap = vm.state.heaps.allocate_with_content(program.as_ref());

src/instruction_handlers/far_call.rs

@@ -10,7 +10,6 @@ use crate::{
     fat_pointer::FatPointer,
     instruction::ExecutionStatus,
     predication::Flags,
-    vm::STORAGE_READ_STORAGE_APPLICATION_CYCLES,
     Instruction, VirtualMachine, World,
 };
 use eravm_stable_interface::{
@@ -19,9 +18,8 @@ use eravm_stable_interface::{
 };
 use u256::U256;
 use zkevm_opcode_defs::{
-    ethereum_types::Address,
     system_params::{EVM_SIMULATOR_STIPEND, MSG_VALUE_SIMULATOR_ADDITIVE_COST},
-    ADDRESS_MSG_VALUE, DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW,
+    ADDRESS_MSG_VALUE,
 };
 
 /// A call to another contract.
@@ -65,16 +63,6 @@ fn far_call<
             };
 
         let failing_part = (|| {
-            let deployer_system_contract_address =
-                Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW as u64);
-            if !vm
-                .world_diff
-                .read_storage_slots
-                .contains(&(deployer_system_contract_address, destination_address))
-            {
-                vm.state.cycle_counts.storage_application_cycles +=
-                    STORAGE_READ_STORAGE_APPLICATION_CYCLES as usize;
-            }
             let decommit_result = vm.world_diff.decommit(
                 world,
                 destination_address,

src/instruction_handlers/storage.rs

@@ -4,7 +4,6 @@ use crate::{
         Arguments, Destination, Register1, Register2, Source, SLOAD_COST, SSTORE_COST,
     },
     instruction::ExecutionStatus,
-    vm::{STORAGE_READ_STORAGE_APPLICATION_CYCLES, STORAGE_WRITE_STORAGE_APPLICATION_CYCLES},
     Instruction, VirtualMachine, World,
 };
 use eravm_stable_interface::{opcodes, Tracer};
@@ -17,14 +16,6 @@ fn sstore<T: Tracer, W: World<T>>(
     instruction_boilerplate::<opcodes::StorageWrite, _, _>(vm, world, tracer, |vm, args, world| {
         let key = Register1::get(args, &mut vm.state);
         let value = Register2::get(args, &mut vm.state);
-        if !vm
-            .world_diff
-            .written_storage_slots
-            .contains(&(vm.state.current_frame.address, key))
-        {
-            vm.state.cycle_counts.storage_application_cycles +=
-                STORAGE_WRITE_STORAGE_APPLICATION_CYCLES as usize;
-        }
 
         let refund = vm
             .world_diff
@@ -62,18 +53,6 @@ fn sload<T: Tracer, W: World<T>>(
     instruction_boilerplate::<opcodes::StorageRead, _, _>(vm, world, tracer, |vm, args, world| {
         let key = Register1::get(args, &mut vm.state);
 
-        if !vm
-            .world_diff
-            .read_storage_slots
-            .contains(&(vm.state.current_frame.address, key))
-            && !vm
-                .world_diff
-                .written_storage_slots
-                .contains(&(vm.state.current_frame.address, key))
-        {
-            vm.state.cycle_counts.storage_application_cycles +=
-                STORAGE_READ_STORAGE_APPLICATION_CYCLES as usize;
-        }
         let (value, refund) =
             vm.world_diff
                 .read_storage(world, vm.state.current_frame.address, key);

src/rollback.rs

@@ -71,18 +71,22 @@ pub struct RollbackableSet<K: Ord> {
 }
 
 impl<T: Ord + Clone> RollbackableSet<T> {
-    /// Adds `key` to the set and returns `true` if it was not already present.
+    /// Adds `key` to the set and returns if it was already present.
     pub fn add(&mut self, key: T) -> bool {
         let is_new = self.map.insert(key.clone(), ()).is_none();
         if is_new {
             self.old_entries.push(key);
         }
-        is_new
+        !is_new
     }
 
     pub fn contains(&self, key: &T) -> bool {
         self.map.contains_key(key)
     }
+
+    pub(crate) fn added_after(&self, snapshot: <Self as Rollback>::Snapshot) -> &[T] {
+        &self.old_entries[snapshot..]
+    }
 }
 
 impl<K: Ord> Rollback for RollbackableSet<K> {

src/state.rs

@@ -204,7 +204,6 @@ pub struct CycleCounts {
     pub sha256_cycles: usize,
     pub secp256v1_verify_cycles: usize,
     pub code_decommitter_cycles: usize,
-    pub storage_application_cycles: usize,
 }
 
 impl PartialEq for CycleCounts {

src/world_diff.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 
 use crate::{
     rollback::{Rollback, RollbackableLog, RollbackableMap, RollbackablePod, RollbackableSet},
+    vm::{STORAGE_READ_STORAGE_APPLICATION_CYCLES, STORAGE_WRITE_STORAGE_APPLICATION_CYCLES},
     StorageInterface,
 };
 use u256::{H160, U256};
@@ -102,14 +103,13 @@ impl WorldDiff {
             .copied()
             .unwrap_or_else(|| world.read_storage(contract, key).unwrap_or_default());
 
-        let refund = if world.is_free_storage_slot(&contract, &key)
-            || self.read_storage_slots.contains(&(contract, key))
+        let refund = if self.read_storage_slots.add((contract, key))
+            || world.is_free_storage_slot(&contract, &key)
         {
             WARM_READ_REFUND
         } else {
             0
         };
-        self.read_storage_slots.add((contract, key));
         self.pubdata_costs.push(0);
         (value, refund)
     }
@@ -131,6 +131,8 @@ impl WorldDiff {
 
         if world.is_free_storage_slot(&contract, &key) {
             self.written_storage_slots.add((contract, key));
+            self.read_storage_slots.add((contract, key));
+
             self.storage_refunds.push(WARM_WRITE_REFUND);
             self.pubdata_costs.push(0);
             return WARM_WRITE_REFUND;
@@ -142,17 +144,12 @@ impl WorldDiff {
             .insert((contract, key), update_cost)
             .unwrap_or(0);
 
-        let refund = if self.written_storage_slots.contains(&(contract, key)) {
+        let refund = if self.written_storage_slots.add((contract, key)) {
             WARM_WRITE_REFUND
+        } else if self.read_storage_slots.add((contract, key)) {
+            COLD_WRITE_AFTER_WARM_READ_REFUND
         } else {
-            self.written_storage_slots.add((contract, key));
-
-            if self.read_storage_slots.contains(&(contract, key)) {
-                COLD_WRITE_AFTER_WARM_READ_REFUND
-            } else {
-                self.read_storage_slots.add((contract, key));
-                0
-            }
+            0
         };
 
         let pubdata_cost = (update_cost as i32) - (prepaid as i32);
@@ -334,6 +331,42 @@ impl WorldDiff {
     pub(crate) fn clear_transient_storage(&mut self) {
         self.transient_storage_changes = Default::default();
     }
+
+    pub fn storage_application_cycles_snapshot(&self) -> StorageApplicationCyclesSnapshot {
+        StorageApplicationCyclesSnapshot {
+            written_storage_slots: self.written_storage_slots.snapshot(),
+            read_storage_slots: self.read_storage_slots.snapshot(),
+            decommitted_hashes: self.decommitted_hashes.snapshot(),
+        }
+    }
+
+    pub fn storage_application_cycles_after(
+        &self,
+        snapshot: &StorageApplicationCyclesSnapshot,
+    ) -> usize {
+        self.written_storage_slots
+            .added_after(snapshot.written_storage_slots)
+            .len()
+            * STORAGE_WRITE_STORAGE_APPLICATION_CYCLES as usize
+            + self
+                .read_storage_slots
+                .added_after(snapshot.read_storage_slots)
+                .iter()
+                .filter(|slot| !self.written_storage_slots.contains(slot))
+                .count()
+                * STORAGE_READ_STORAGE_APPLICATION_CYCLES as usize
+            + self
+                .decommitted_hashes
+                .added_after(snapshot.decommitted_hashes)
+                .len()
+                * STORAGE_READ_STORAGE_APPLICATION_CYCLES as usize
+    }
+}
+
+pub struct StorageApplicationCyclesSnapshot {
+    written_storage_slots: <RollbackableSet<(H160, U256)> as Rollback>::Snapshot,
+    read_storage_slots: <RollbackableSet<(H160, U256)> as Rollback>::Snapshot,
+    decommitted_hashes: <RollbackableSet<U256> as Rollback>::Snapshot,
 }
 
 #[derive(Clone, PartialEq, Debug)]