feat(cheatcodes): additional cheatcodes to aid in symbolic testing (#8807)

* feat(cheatcodes): additional cheatcodes to aid in symbolic testing

* Support copies from arbitrary storage, docs

* Changes after review:
- separate cheatcodes tests with specific seed
- better way to match mocked function
- arbitrary_storage_end instead multiple calls
- generate arbitrary value only when needed

* Update crates/cheatcodes/src/utils.rs

Co-authored-by: DaniPopes <[email protected]>

* Fix tests with isolate-by-default

---------

Co-authored-by: DaniPopes <[email protected]>

Author: grandizzy

crates/cheatcodes/assets/cheatcodes.json

@@ -473,6 +473,15 @@ interface Vm {
     function mockCallRevert(address callee, uint256 msgValue, bytes calldata data, bytes calldata revertData)
         external;
 
+    /// Whenever a call is made to `callee` with calldata `data`, this cheatcode instead calls
+    /// `target` with the same calldata. This functionality is similar to a delegate call made to
+    /// `target` contract from `callee`.
+    /// Can be used to substitute a call to a function with another implementation that captures
+    /// the primary logic of the original function but is easier to reason about.
+    /// If calldata is not a strict match then partial match by selector is attempted.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function mockFunction(address callee, address target, bytes calldata data) external;
+
     // --- Impersonation (pranks) ---
 
     /// Sets the *next* call's `msg.sender` to be the input address.
@@ -2303,6 +2312,14 @@ interface Vm {
     /// Unpauses collection of call traces.
     #[cheatcode(group = Utilities)]
     function resumeTracing() external view;
+
+    /// Utility cheatcode to copy storage of `from` contract to another `to` contract.
+    #[cheatcode(group = Utilities)]
+    function copyStorage(address from, address to) external;
+
+    /// Utility cheatcode to set arbitrary storage for given target address.
+    #[cheatcode(group = Utilities)]
+    function setArbitraryStorage(address target) external;
 }
 }
 

crates/cheatcodes/spec/src/vm.rs

@@ -13,6 +13,7 @@ use foundry_evm_core::{
     backend::{DatabaseExt, RevertSnapshotAction},
     constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS},
 };
+use rand::Rng;
 use revm::{
     primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY},
     InnerEvmContext,
@@ -89,7 +90,25 @@ impl Cheatcode for loadCall {
         let Self { target, slot } = *self;
         ensure_not_precompile!(&target, ccx);
         ccx.ecx.load_account(target)?;
-        let val = ccx.ecx.sload(target, slot.into())?;
+        let mut val = ccx.ecx.sload(target, slot.into())?;
+
+        if val.is_cold && val.data.is_zero() {
+            if ccx.state.arbitrary_storage.is_arbitrary(&target) {
+                // If storage slot is untouched and load from a target with arbitrary storage,
+                // then set random value for current slot.
+                let rand_value = ccx.state.rng().gen();
+                ccx.state.arbitrary_storage.save(ccx.ecx, target, slot.into(), rand_value);
+                val.data = rand_value;
+            } else if ccx.state.arbitrary_storage.is_copy(&target) {
+                // If storage slot is untouched and load from a target that copies storage from
+                // a source address with arbitrary storage, then copy existing arbitrary value.
+                // If no arbitrary value generated yet, then the random one is saved and set.
+                let rand_value = ccx.state.rng().gen();
+                val.data =
+                    ccx.state.arbitrary_storage.copy(ccx.ecx, target, slot.into(), rand_value);
+            }
+        }
+
         Ok(val.abi_encode())
     }
 }

crates/cheatcodes/src/evm.rs

@@ -89,6 +89,15 @@ impl Cheatcode for mockCallRevert_1Call {
     }
 }
 
+impl Cheatcode for mockFunctionCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { callee, target, data } = self;
+        state.mocked_functions.entry(*callee).or_default().insert(data.clone(), *target);
+
+        Ok(Default::default())
+    }
+}
+
 #[allow(clippy::ptr_arg)] // Not public API, doesn't matter
 fn mock_call(
     state: &mut Cheatcodes,

crates/cheatcodes/src/evm/mock.rs

@@ -41,7 +41,7 @@ use revm::{
         EOFCreateInputs, EOFCreateKind, Gas, InstructionResult, Interpreter, InterpreterAction,
         InterpreterResult,
     },
-    primitives::{BlockEnv, CreateScheme, EVMError, SpecId, EOF_MAGIC_BYTES},
+    primitives::{BlockEnv, CreateScheme, EVMError, EvmStorageSlot, SpecId, EOF_MAGIC_BYTES},
     EvmContext, InnerEvmContext, Inspector,
 };
 use rustc_hash::FxHashMap;
@@ -254,6 +254,89 @@ impl GasMetering {
     }
 }
 
+/// Holds data about arbitrary storage.
+#[derive(Clone, Debug, Default)]
+pub struct ArbitraryStorage {
+    /// Mapping of arbitrary storage addresses to generated values (slot, arbitrary value).
+    /// (SLOADs return random value if storage slot wasn't accessed).
+    /// Changed values are recorded and used to copy storage to different addresses.
+    pub values: HashMap<Address, HashMap<U256, U256>>,
+    /// Mapping of address with storage copied to arbitrary storage address source.
+    pub copies: HashMap<Address, Address>,
+}
+
+impl ArbitraryStorage {
+    /// Whether the given address has arbitrary storage.
+    pub fn is_arbitrary(&self, address: &Address) -> bool {
+        self.values.contains_key(address)
+    }
+
+    /// Whether the given address is a copy of an address with arbitrary storage.
+    pub fn is_copy(&self, address: &Address) -> bool {
+        self.copies.contains_key(address)
+    }
+
+    /// Marks an address with arbitrary storage.
+    pub fn mark_arbitrary(&mut self, address: &Address) {
+        self.values.insert(*address, HashMap::default());
+    }
+
+    /// Maps an address that copies storage with the arbitrary storage address.
+    pub fn mark_copy(&mut self, from: &Address, to: &Address) {
+        if self.is_arbitrary(from) {
+            self.copies.insert(*to, *from);
+        }
+    }
+
+    /// Saves arbitrary storage value for a given address:
+    /// - store value in changed values cache.
+    /// - update account's storage with given value.
+    pub fn save<DB: DatabaseExt>(
+        &mut self,
+        ecx: &mut InnerEvmContext<DB>,
+        address: Address,
+        slot: U256,
+        data: U256,
+    ) {
+        self.values.get_mut(&address).expect("missing arbitrary address entry").insert(slot, data);
+        if let Ok(mut account) = ecx.load_account(address) {
+            account.storage.insert(slot, EvmStorageSlot::new(data));
+        }
+    }
+
+    /// Copies arbitrary storage value from source address to the given target address:
+    /// - if a value is present in arbitrary values cache, then update target storage and return
+    ///   existing value.
+    /// - if no value was yet generated for given slot, then save new value in cache and update both
+    ///   source and target storages.
+    pub fn copy<DB: DatabaseExt>(
+        &mut self,
+        ecx: &mut InnerEvmContext<DB>,
+        target: Address,
+        slot: U256,
+        new_value: U256,
+    ) -> U256 {
+        let source = self.copies.get(&target).expect("missing arbitrary copy target entry");
+        let storage_cache = self.values.get_mut(source).expect("missing arbitrary source storage");
+        let value = match storage_cache.get(&slot) {
+            Some(value) => *value,
+            None => {
+                storage_cache.insert(slot, new_value);
+                // Update source storage with new value.
+                if let Ok(mut source_account) = ecx.load_account(*source) {
+                    source_account.storage.insert(slot, EvmStorageSlot::new(new_value));
+                }
+                new_value
+            }
+        };
+        // Update target storage with new value.
+        if let Ok(mut target_account) = ecx.load_account(target) {
+            target_account.storage.insert(slot, EvmStorageSlot::new(value));
+        }
+        value
+    }
+}
+
 /// List of transactions that can be broadcasted.
 pub type BroadcastableTransactions = VecDeque<BroadcastableTransaction>;
 
@@ -320,6 +403,9 @@ pub struct Cheatcodes {
     // **Note**: inner must a BTreeMap because of special `Ord` impl for `MockCallDataContext`
     pub mocked_calls: HashMap<Address, BTreeMap<MockCallDataContext, MockCallReturnData>>,
 
+    /// Mocked functions. Maps target address to be mocked to pair of (calldata, mock address).
+    pub mocked_functions: HashMap<Address, HashMap<Bytes, Address>>,
+
     /// Expected calls
     pub expected_calls: ExpectedCallTracker,
     /// Expected emits
@@ -368,6 +454,9 @@ pub struct Cheatcodes {
 
     /// Ignored traces.
     pub ignored_traces: IgnoredTraces,
+
+    /// Addresses with arbitrary storage.
+    pub arbitrary_storage: ArbitraryStorage,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -396,6 +485,7 @@ impl Cheatcodes {
             recorded_account_diffs_stack: Default::default(),
             recorded_logs: Default::default(),
             mocked_calls: Default::default(),
+            mocked_functions: Default::default(),
             expected_calls: Default::default(),
             expected_emits: Default::default(),
             allowed_mem_writes: Default::default(),
@@ -410,6 +500,7 @@ impl Cheatcodes {
             breakpoints: Default::default(),
             rng: Default::default(),
             ignored_traces: Default::default(),
+            arbitrary_storage: Default::default(),
         }
     }
 
@@ -1045,14 +1136,22 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
     }
 
     #[inline]
-    fn step_end(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext<DB>) {
+    fn step_end(&mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext<DB>) {
         if self.gas_metering.paused {
             self.meter_gas_end(interpreter);
         }
 
         if self.gas_metering.touched {
             self.meter_gas_check(interpreter);
         }
+
+        // `setArbitraryStorage` and `copyStorage`: add arbitrary values to storage.
+        if (self.arbitrary_storage.is_arbitrary(&interpreter.contract().target_address) ||
+            self.arbitrary_storage.is_copy(&interpreter.contract().target_address)) &&
+            interpreter.current_opcode() == op::SLOAD
+        {
+            self.arbitrary_storage_end(interpreter, ecx);
+        }
     }
 
     fn log(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext<DB>, log: &Log) {
@@ -1465,6 +1564,43 @@ impl Cheatcodes {
         }
     }
 
+    /// Generates or copies arbitrary values for storage slots.
+    /// Invoked in inspector `step_end` (when the current opcode is not executed), if current opcode
+    /// to execute is `SLOAD` and storage slot is cold.
+    /// Ensures that in next step (when `SLOAD` opcode is executed) an arbitrary value is returned:
+    /// - copies the existing arbitrary storage value (or the new generated one if no value in
+    ///   cache) from mapped source address to the target address.
+    /// - generates arbitrary value and saves it in target address storage.
+    #[cold]
+    fn arbitrary_storage_end<DB: DatabaseExt>(
+        &mut self,
+        interpreter: &mut Interpreter,
+        ecx: &mut EvmContext<DB>,
+    ) {
+        let key = try_or_return!(interpreter.stack().peek(0));
+        let target_address = interpreter.contract().target_address;
+        if let Ok(value) = ecx.sload(target_address, key) {
+            if value.is_cold && value.data.is_zero() {
+                let arbitrary_value = self.rng().gen();
+                if self.arbitrary_storage.is_copy(&target_address) {
+                    self.arbitrary_storage.copy(
+                        &mut ecx.inner,
+                        target_address,
+                        key,
+                        arbitrary_value,
+                    );
+                } else {
+                    self.arbitrary_storage.save(
+                        &mut ecx.inner,
+                        target_address,
+                        key,
+                        arbitrary_value,
+                    );
+                }
+            }
+        }
+    }
+
     /// Records storage slots reads and writes.
     #[cold]
     fn record_accesses(&mut self, interpreter: &mut Interpreter) {