caller decides what storage slots are free

src/instruction_handlers/storage.rs

@@ -12,27 +12,22 @@ use crate::{
 
 fn sstore(vm: &mut VirtualMachine, instruction: *const Instruction) -> InstructionResult {
     instruction_boilerplate_with_panic(vm, instruction, |vm, args, continue_normally| {
-        let key = Register1::get(args, &mut vm.state);
-        let value = Register2::get(args, &mut vm.state);
-        let address = vm.state.current_frame.address;
-
-        let update_cost = vm.world.world.cost_of_writing_storage(address, key, value);
-        let prepaid = vm.world.prepaid_for_write(address, key);
-
-        // Note, that the diff may be negative, e.g. in case the new write returns to the previous value.
-        vm.state.current_frame.total_pubdata_spent += (update_cost as i32) - (prepaid as i32);
-
-        vm.world.insert_prepaid_for_write(address, key, update_cost);
-
         if vm.state.current_frame.is_static {
             return Ok(&PANIC);
         }
 
-        let refund = vm.world.write_storage(address, key, value);
+        let key = Register1::get(args, &mut vm.state);
+        let value = Register2::get(args, &mut vm.state);
+
+        let (refund, pubdata_change) =
+            vm.world
+                .write_storage(vm.state.current_frame.address, key, value);
 
         assert!(refund <= SSTORE_COST);
         vm.state.current_frame.gas += refund;
 
+        vm.state.current_frame.total_pubdata_spent += pubdata_change;
+
         continue_normally
     })
 }

src/lib.rs

@@ -36,5 +36,9 @@ pub trait World {
     /// There is no write_storage; [ModifiedWorld::get_storage_changes] gives a list of all storage changes.
     fn read_storage(&mut self, contract: H160, key: U256) -> U256;
 
+    /// Computes the cost of writing a storage slot.
     fn cost_of_writing_storage(&mut self, contract: H160, key: U256, new_value: U256) -> u32;
+
+    /// Returns if the storage slot is free both in terms of gas and pubdata.
+    fn is_free_storage_slot(&self, contract: &H160, key: &U256) -> bool;
 }

src/modified_world.rs

@@ -4,14 +4,10 @@ use crate::{
     rollback::{Rollback, RollbackableLog, RollbackableMap, RollbackableSet},
     World,
 };
-use u256::{H160, H256, U256};
-use zkevm_opcode_defs::{
-    blake2::Blake2s256,
-    sha3::Digest,
-    system_params::{
-        STORAGE_ACCESS_COLD_READ_COST, STORAGE_ACCESS_COLD_WRITE_COST,
-        STORAGE_ACCESS_WARM_READ_COST, STORAGE_ACCESS_WARM_WRITE_COST,
-    },
+use u256::{H160, U256};
+use zkevm_opcode_defs::system_params::{
+    STORAGE_ACCESS_COLD_READ_COST, STORAGE_ACCESS_COLD_WRITE_COST, STORAGE_ACCESS_WARM_READ_COST,
+    STORAGE_ACCESS_WARM_WRITE_COST,
 };
 
 /// The global state including pending modifications that are written only at
@@ -73,15 +69,15 @@ impl ModifiedWorld {
     }
 
     /// Returns the storage slot's value and a refund based on its hot/cold status.
-    pub fn read_storage(&mut self, contract: H160, key: U256) -> (U256, u32) {
+    pub(crate) fn read_storage(&mut self, contract: H160, key: U256) -> (U256, u32) {
         let value = self
             .storage_changes
             .as_ref()
             .get(&(contract, key))
             .cloned()
             .unwrap_or_else(|| self.world.read_storage(contract, key));
 
-        let refund = if is_storage_key_free(&contract, &key)
+        let refund = if self.world.is_free_storage_slot(&contract, &key)
             || self.read_storage_slots.contains(&(contract, key))
         {
             WARM_READ_REFUND
@@ -93,15 +89,24 @@ impl ModifiedWorld {
         (value, refund)
     }
 
-    /// Returns the refund based the hot/cold status of the storage slot.
-    pub fn write_storage(&mut self, contract: H160, key: U256, value: U256) -> u32 {
+    /// Returns the refund based the hot/cold status of the storage slot and the change in pubdata.
+    pub(crate) fn write_storage(&mut self, contract: H160, key: U256, value: U256) -> (u32, i32) {
         self.storage_changes.insert((contract, key), value);
 
-        if is_storage_key_free(&contract, &key)
-            || self
-                .written_storage_slots
-                .as_ref()
-                .contains_key(&(contract, key))
+        if self.world.is_free_storage_slot(&contract, &key) {
+            return (WARM_WRITE_REFUND, 0);
+        }
+
+        let update_cost = self.world.cost_of_writing_storage(contract, key, value);
+        let prepaid = self
+            .paid_changes
+            .insert((contract, key), update_cost)
+            .unwrap_or(0);
+
+        let refund = if self
+            .written_storage_slots
+            .as_ref()
+            .contains_key(&(contract, key))
         {
             WARM_WRITE_REFUND
         } else {
@@ -113,19 +118,9 @@ impl ModifiedWorld {
                 self.read_storage_slots.add((contract, key));
                 0
             }
-        }
-    }
-
-    pub fn prepaid_for_write(&self, address: H160, key: U256) -> u32 {
-        self.paid_changes
-            .as_ref()
-            .get(&(address, key))
-            .cloned()
-            .unwrap_or(0u32)
-    }
+        };
 
-    pub fn insert_prepaid_for_write(&mut self, address: H160, key: U256, price: u32) {
-        self.paid_changes.insert((address, key), price)
+        (refund, (update_cost as i32) - (prepaid as i32))
     }
 
     pub fn get_storage_changes(&self) -> &BTreeMap<(H160, U256), U256> {
@@ -208,42 +203,3 @@ pub(crate) type Snapshot = (
 const WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST - STORAGE_ACCESS_WARM_READ_COST;
 const WARM_WRITE_REFUND: u32 = STORAGE_ACCESS_COLD_WRITE_COST - STORAGE_ACCESS_WARM_WRITE_COST;
 const COLD_WRITE_AFTER_WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST;
-
-pub(crate) fn is_storage_key_free(address: &H160, key: &U256) -> bool {
-    let storage_key_for_eth_balance = U256([
-        4209092924407300373,
-        6927221427678996148,
-        4194905989268492595,
-        15931007429432312239,
-    ]);
-    if address == &SYSTEM_CONTEXT_ADDRESS {
-        return true;
-    }
-
-    let keyy = U256::from_little_endian(&raw_hashed_key(address, &u256_to_h256(*key)));
-
-    if keyy == storage_key_for_eth_balance {
-        return true;
-    }
-
-    false
-}
-
-pub const SYSTEM_CONTEXT_ADDRESS: H160 = H160([
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0x00, 0x00, 0x80, 0x0b,
-]);
-
-pub fn u256_to_h256(num: U256) -> H256 {
-    let mut bytes = [0u8; 32];
-    num.to_big_endian(&mut bytes);
-    H256::from_slice(&bytes)
-}
-
-fn raw_hashed_key(address: &H160, key: &H256) -> [u8; 32] {
-    let mut bytes = [0u8; 64];
-    bytes[12..32].copy_from_slice(&address.0);
-    U256::from(key.to_fixed_bytes()).to_big_endian(&mut bytes[32..64]);
-
-    Blake2s256::digest(bytes).into()
-}

src/rollback.rs

@@ -14,10 +14,11 @@ pub struct RollbackableMap<K: Ord, V> {
     old_entries: Vec<(K, Option<V>)>,
 }
 
-impl<K: Ord + Clone, V> RollbackableMap<K, V> {
-    pub fn insert(&mut self, key: K, value: V) {
-        self.old_entries
-            .push((key.clone(), self.map.insert(key, value)));
+impl<K: Ord + Clone, V: Clone> RollbackableMap<K, V> {
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        let old_value = self.map.insert(key.clone(), value);
+        self.old_entries.push((key.clone(), old_value.clone()));
+        old_value
     }
 }
 

src/testworld.rs

@@ -70,4 +70,8 @@ impl World for TestWorld {
     ) -> u32 {
         50
     }
+
+    fn is_free_storage_slot(&self, _contract: &u256::H160, _key: &U256) -> bool {
+        false
+    }
 }