free addrs should cost 0

src/decommit.rs

@@ -1,5 +1,5 @@
 use crate::{modified_world::ModifiedWorld, program::Program, World};
-use u256::{H160, U256};
+use u256::{H160, H256, U256};
 use zkevm_opcode_defs::{
     ethereum_types::Address, system_params::DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW,
 };
@@ -102,6 +102,14 @@ pub(crate) fn u256_into_address(source: U256) -> H160 {
     result
 }
 
+pub(crate) fn u256_into_h256(source: U256) -> H256 {
+    let mut result = H256::zero();
+    let mut bytes = [0; 32];
+    source.to_big_endian(&mut bytes);
+    result.assign_from_slice(&bytes[..]);
+    result
+}
+
 pub(crate) fn is_kernel(address: U256) -> bool {
     address < (1 << 16).into()
 }

src/instruction_handlers/storage.rs

@@ -1,4 +1,5 @@
-use u256::U256;
+use u256::{H160, H256, U256};
+use zkevm_opcode_defs::{blake2::Blake2s256, sha3::Digest};
 
 use super::{
     common::{instruction_boilerplate, instruction_boilerplate_with_panic},
@@ -43,24 +44,85 @@ fn get_pubdata_price_bytes(initial_value: U256, final_value: U256, is_initial: b
     }
 }
 
+fn base_price_for_write_query(vm: &mut VirtualMachine, key: U256, new_value: U256) -> u32 {
+    let contract = vm.state.current_frame.address;
+    let initial_value = vm
+        .world
+        .get_initial_value(&contract, &key)
+        .unwrap_or(vm.world.world.read_storage(contract, key));
+
+    let is_initial = vm
+        .world
+        .is_write_initial(vm.state.current_frame.address, key);
+
+    if is_storage_key_free(&vm.state.current_frame.address, &key) || initial_value == new_value {
+        return 0;
+    }
+
+    let ret = get_pubdata_price_bytes(initial_value, new_value, is_initial);
+    ret
+}
+
+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,
+]);
+
+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()
+}
+
+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
+}
+
 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 is_initial = vm
-            .world
-            .is_write_initial(vm.state.current_frame.address, key);
-
-        let (initial_value, _) = vm.world.read_storage(vm.state.current_frame.address, key);
+        vm.world
+            .set_initial_value(vm.state.current_frame.address, key, value);
 
-        let to_pay_by_user = get_pubdata_price_bytes(initial_value, value, is_initial);
+        let to_pay_by_user = base_price_for_write_query(vm, key, value);
         let prepaid = vm
             .world
             .prepaid_for_write(vm.state.current_frame.address, key);
 
         // Note, that the diff may be negative, e.g. in case the new write returns to the previous value.
         let diff = (to_pay_by_user as i32) - (prepaid as i32);
+        println!(
+            ">>> StorageKey {{ account: AccountTreeId {{ address: {:?} }}, key: {:?} }} {:?} {to_pay_by_user} {prepaid}",
+            vm.state.current_frame.address,
+            u256_to_h256(key),
+            u256_to_h256(value)
+        );
         vm.state.current_frame.total_pubdata_spent += diff;
 
         vm.world
@@ -86,6 +148,9 @@ fn sload(vm: &mut VirtualMachine, instruction: *const Instruction) -> Instructio
         let key = Register1::get(args, &mut vm.state);
         let (value, refund) = vm.world.read_storage(vm.state.current_frame.address, key);
 
+        vm.world
+            .set_initial_value(vm.state.current_frame.address, key, value);
+
         assert!(refund <= SLOAD_COST);
         vm.state.current_frame.gas += refund;
 

src/modified_world.rs

@@ -20,6 +20,7 @@ pub struct ModifiedWorld {
     events: RollbackableLog<Event>,
     l2_to_l1_logs: RollbackableLog<L2ToL1Log>,
     paid_changes: RollbackableMap<(H160, U256), u32>,
+    initial_values: RollbackableMap<(H160, U256), U256>,
 
     // The fields below are only rolled back when the whole VM is rolled back.
     pub(crate) decommitted_hashes: RollbackableSet<U256>,
@@ -65,6 +66,7 @@ impl ModifiedWorld {
             read_storage_slots: Default::default(),
             written_storage_slots: Default::default(),
             paid_changes: Default::default(),
+            initial_values: Default::default(),
         }
     }
 
@@ -121,6 +123,16 @@ impl ModifiedWorld {
         self.paid_changes.insert((address, key), price)
     }
 
+    pub fn set_initial_value(&mut self, address: H160, key: U256, value: U256) {
+        if !self.initial_values.as_ref().contains_key(&(address, key)) {
+            self.initial_values.insert((address, key), value);
+        }
+    }
+
+    pub fn get_initial_value(&self, address: &H160, key: &U256) -> Option<U256> {
+        self.initial_values.as_ref().get(&(*address, *key)).copied()
+    }
+
     pub fn get_storage_changes(&self) -> &BTreeMap<(H160, U256), U256> {
         self.storage_changes.as_ref()
     }

src/vm.rs

@@ -14,6 +14,7 @@ use zkevm_opcode_defs::system_params::NEW_FRAME_MEMORY_STIPEND;
 pub struct Settings {
     pub default_aa_code_hash: [u8; 32],
     pub evm_interpreter_code_hash: [u8; 32],
+    pub storage_key_for_eth_balance: [u8; 32],
 
     /// Writing to this address in the bootloader's heap suspends execution
     pub hook_address: u32,