fix(vm.cool) Persist storage changes

Author: ditto-eth

crates/cheatcodes/defs/src/vm.rs

@@ -255,7 +255,8 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Unsafe)]
     function store(address target, bytes32 slot, bytes32 value) external;
 
-    /// Marks the slots of an account and the account address as cold.
+    /// Marks the `target` address cold, and is a no-op if the address is already cold.
+    /// All storage slots are also made cold, but their values are preserved.
     #[cheatcode(group = Evm, safety = Unsafe)]
     function cool(address target) external;
 

crates/cheatcodes/src/evm.rs

@@ -1,6 +1,7 @@
 //! Implementations of [`Evm`](crate::Group::Evm) cheatcodes.
 
-use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
+use crate::{inspector::AddressState, Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
+
 use alloy_primitives::{Address, Bytes, U256};
 use alloy_sol_types::SolValue;
 use ethers_core::utils::{Genesis, GenesisAccount};
@@ -9,7 +10,7 @@ use foundry_common::fs::read_json_file;
 use foundry_evm_core::backend::DatabaseExt;
 use foundry_utils::types::ToAlloy;
 use revm::{
-    primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY},
+    primitives::{Account, Bytecode, HashMap as rHashMap, SpecId, KECCAK_EMPTY},
     EVMData,
 };
 use std::{collections::HashMap, path::Path};
@@ -302,11 +303,11 @@ impl Cheatcode for storeCall {
 
 impl Cheatcode for coolCall {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
-        let Self { target } = self;
-        if let Some(account) = ccx.data.journaled_state.state.get_mut(target) {
-            account.unmark_touch();
-            account.storage.clear();
-        }
+        let Self { target } = *self;
+        ensure_not_precompile!(&target, ccx);
+        // TODO: prevent or warn about cooling the to/from address in a tx
+        ccx.state.addresses.insert(target, AddressState::Cool);
+        ccx.state.address_storage.insert(target, rHashMap::new());
         Ok(Default::default())
     }
 }

crates/cheatcodes/src/inspector.rs

@@ -29,7 +29,7 @@ use foundry_utils::types::ToEthers;
 use itertools::Itertools;
 use revm::{
     interpreter::{opcode, CallInputs, CreateInputs, Gas, InstructionResult, Interpreter},
-    primitives::{BlockEnv, CreateScheme, TransactTo},
+    primitives::{BlockEnv, CreateScheme, HashMap as rHashMap, TransactTo},
     EVMData, Inspector,
 };
 use serde_json::Value;
@@ -200,6 +200,174 @@ pub struct Cheatcodes {
     /// Breakpoints supplied by the `breakpoint` cheatcode.
     /// `char -> (address, pc)`
     pub breakpoints: Breakpoints,
+
+    /// Track if cool cheatcode was called on each address
+    pub addresses: rHashMap<Address, AddressState>,
+    /// Track if an addresses storage slot is cool or not
+    pub address_storage: rHashMap<Address, rHashMap<U256, StorageSlotState>>,
+    /// How much gas to charge in the next step (op code) based on cool cheatcode calculations
+    pub additional_gas_next_op: u64,
+}
+
+/// Whether an Address is accessed or not
+#[derive(Clone, PartialEq, Debug)]
+pub enum AddressState {
+    /// if already accessed, then charge WARM_STORAGE_READ_COST (100)
+    Warm,
+    /// charge COLD_ACCOUNT_ACCESS_COST (2600)
+    Cool,
+}
+
+/// Whether a Storage Slot is warm or already been modified
+#[derive(Clone, PartialEq, Debug)]
+pub enum StorageSlotState {
+    /// charge extra based on SSTORE calculations
+    WarmWithSLOAD,
+    /// if SSTORE already happened, don't charge extra
+    WarmWithSSTORE,
+    /// same as if empty
+    Cool,
+}
+
+/// Function to charge extra gas per opcode based on cool cheatcode
+fn add_gas_from_cool_cheatcode<DB: DatabaseExt>(
+    state: &mut Cheatcodes,
+    interpreter: &mut Interpreter,
+    data: &mut EVMData<'_, DB>,
+) -> InstructionResult {
+    // For gas costs, see https://eips.ethereum.org/EIPS/eip-2200, https://eips.ethereum.org/EIPS/eip-2929
+
+    // if previous step added gas, add it once
+    // note that all the opcodes will already have a cost (usually 100)
+    // so adding until it hits the gas expected for a cold key/address
+    if state.additional_gas_next_op > 0 {
+        interpreter.gas.record_cost(state.additional_gas_next_op);
+        state.additional_gas_next_op = 0;
+    }
+
+    // if cool cheatcode was ever called on this address
+    let contract_address = interpreter.contract().address;
+
+    if state.addresses.get(&contract_address).is_some() {
+        if let Some(contract_storage) = state.address_storage.get_mut(&contract_address) {
+            match interpreter.current_opcode() {
+                // via AccessListTracer
+                opcode::EXTCODECOPY |
+                opcode::EXTCODEHASH |
+                opcode::EXTCODESIZE |
+                opcode::BALANCE |
+                opcode::SELFDESTRUCT => {
+                    // address is first parameter
+                    if let Ok(slot) = interpreter.stack().peek(0) {
+                        let addr: Address = Address::from_word(slot.into());
+
+                        // COLD_ACCOUNT_ACCESS_COST is 2600
+                        // check this is done once per address, unless cheatcode is called again
+                        // ignore if same as contract address
+                        if addr != contract_address &&
+                            state.addresses.get(&addr) == Some(&AddressState::Cool)
+                        {
+                            state.additional_gas_next_op = 2500;
+                            state.addresses.insert(addr, AddressState::Warm);
+                        }
+                    }
+                }
+                // via AccessListTracer
+                opcode::DELEGATECALL | opcode::CALL | opcode::STATICCALL | opcode::CALLCODE => {
+                    // address is second parameter
+                    if let Ok(slot) = interpreter.stack().peek(1) {
+                        let addr: Address = Address::from_word(slot.into());
+
+                        // COLD_ACCOUNT_ACCESS_COST is 2600
+                        // check this is done once per address, unless cheatcode is called again
+                        // ignore if same as contract address
+                        if addr != contract_address &&
+                            state.addresses.get(&addr) == Some(&AddressState::Cool)
+                        {
+                            state.additional_gas_next_op = 2500;
+                            state.addresses.insert(addr, AddressState::Warm);
+                        }
+                    }
+                }
+                opcode::SLOAD => {
+                    let key = try_or_continue!(interpreter.stack().peek(0));
+
+                    let account = data.journaled_state.state().get(&contract_address).unwrap();
+                    if account.storage.get(&key).is_some() {
+                        match contract_storage.get(&key) {
+                            None | Some(StorageSlotState::Cool) => {
+                                // COLD_SLOAD_COST is 2100
+                                state.additional_gas_next_op = 2000;
+                                contract_storage.insert(key, StorageSlotState::WarmWithSLOAD);
+                            }
+                            Some(_) => {}
+                        }
+                    } else {
+                        contract_storage.insert(key, StorageSlotState::WarmWithSLOAD);
+                    }
+                }
+                opcode::SSTORE => {
+                    let key = try_or_continue!(interpreter.stack().peek(0));
+                    let val = try_or_continue!(interpreter.stack().peek(1));
+
+                    let account = data.journaled_state.state().get(&contract_address).unwrap();
+                    if account.storage.get(&key).is_some() {
+                        // only add gas the first time the storage is touched again
+                        match contract_storage.get(&key) {
+                            Some(StorageSlotState::WarmWithSLOAD) => {
+                                // cool keeps the slot value changes
+                                // as if the previous_or_original_value = present_value`
+                                // so include the extra gas
+                                let slot = account.storage.get(&key).unwrap();
+                                if val != slot.present_value &&
+                                    slot.present_value != slot.previous_or_original_value
+                                {
+                                    if slot.present_value == U256::ZERO {
+                                        // SSTORE_SET_GAS is 20000
+                                        state.additional_gas_next_op += 20000 - 100
+                                    } else {
+                                        // SSTORE_RESET_GAS is 5000 - COLD_SLOAD_COST (2100)
+                                        state.additional_gas_next_op += 2900 - 100
+                                    }
+                                }
+
+                                // set slot is_warm to true
+                                contract_storage.insert(key, StorageSlotState::WarmWithSSTORE);
+                            }
+                            None | Some(StorageSlotState::Cool) => {
+                                // Means SSTORE was called without SLOAD before
+                                // COLD_SLOAD_COST is 2100
+                                state.additional_gas_next_op = 2100;
+
+                                // cool keeps the slot value changes
+                                // as if the previous_or_original_value = present_value`
+                                // so include the extra gas
+                                let slot = account.storage.get(&key).unwrap();
+                                if val != slot.present_value &&
+                                    slot.present_value != slot.previous_or_original_value
+                                {
+                                    if slot.present_value == U256::ZERO {
+                                        // SSTORE_SET_GAS is 20000
+                                        state.additional_gas_next_op += 20000 - 100
+                                    } else {
+                                        // SSTORE_RESET_GAS is 5000 - COLD_SLOAD_COST (2100)
+                                        state.additional_gas_next_op += 2900 - 100
+                                    }
+                                }
+                                contract_storage.insert(key, StorageSlotState::WarmWithSSTORE);
+                            }
+                            Some(StorageSlotState::WarmWithSSTORE) => {}
+                        }
+                    } else {
+                        contract_storage.insert(key, StorageSlotState::WarmWithSSTORE);
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
+
+    InstructionResult::Continue
 }
 
 impl Cheatcodes {
@@ -523,6 +691,16 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
         InstructionResult::Continue
     }
 
+    fn step_end(
+        &mut self,
+        interpreter: &mut Interpreter,
+        data: &mut EVMData<'_, DB>,
+        eval: InstructionResult,
+    ) -> InstructionResult {
+        add_gas_from_cool_cheatcode(self, interpreter, data);
+        eval
+    }
+
     fn log(&mut self, _: &mut EVMData<'_, DB>, address: &Address, topics: &[B256], data: &Bytes) {
         if !self.expected_emits.is_empty() {
             expect::handle_expect_emit(self, address, topics, data);