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

crates/cheatcodes/spec/src/vm.rs

@@ -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/src/evm.rs

@@ -13,8 +13,9 @@ 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},
+    primitives::{Account, Bytecode, EvmStorageSlot, SpecId, KECCAK_EMPTY},
     InnerEvmContext,
 };
 use std::{
@@ -89,7 +90,13 @@ 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())?;
+        // Generate random value if target should have arbitrary storage and storage slot untouched.
+        if ccx.state.arbitrary_storage.contains(&target) && val.is_cold && val.data.is_zero() {
+            val.data = ccx.state.rng().gen();
+            let mut account = ccx.ecx.load_account(target)?;
+            account.storage.insert(slot.into(), EvmStorageSlot::new(val.data));
+        }
         Ok(val.abi_encode())
     }
 }

crates/cheatcodes/src/evm/mock.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/inspector.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;
@@ -320,6 +320,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 +371,10 @@ pub struct Cheatcodes {
 
     /// Ignored traces.
     pub ignored_traces: IgnoredTraces,
+
+    /// Addresses that should have arbitrary storage generated (SLOADs return random value if
+    /// storage slot wasn't accessed).
+    pub arbitrary_storage: Vec<Address>,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -396,6 +403,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 +418,7 @@ impl Cheatcodes {
             breakpoints: Default::default(),
             rng: Default::default(),
             ignored_traces: Default::default(),
+            arbitrary_storage: Default::default(),
         }
     }
 
@@ -1045,14 +1054,18 @@ 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);
         }
+
+        if self.arbitrary_storage.contains(&interpreter.contract().target_address) {
+            self.ensure_arbitrary_storage(interpreter, ecx);
+        }
     }
 
     fn log(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext<DB>, log: &Log) {
@@ -1465,6 +1478,26 @@ impl Cheatcodes {
         }
     }
 
+    /// Generates arbitrary values for storage slots.
+    #[cold]
+    fn ensure_arbitrary_storage<DB: DatabaseExt>(
+        &mut self,
+        interpreter: &mut Interpreter,
+        ecx: &mut EvmContext<DB>,
+    ) {
+        if interpreter.current_opcode() == op::SLOAD {
+            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() {
+                    if let Ok(mut target_account) = ecx.load_account(target_address) {
+                        target_account.storage.insert(key, EvmStorageSlot::new(self.rng().gen()));
+                    }
+                }
+            }
+        }
+    }
+
     /// Records storage slots reads and writes.
     #[cold]
     fn record_accesses(&mut self, interpreter: &mut Interpreter) {

crates/cheatcodes/src/utils.rs

@@ -1,6 +1,6 @@
 //! Implementations of [`Utilities`](spec::Group::Utilities) cheatcodes.
 
-use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
+use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
 use alloy_primitives::{Address, U256};
 use alloy_sol_types::SolValue;
 use foundry_common::ens::namehash;
@@ -149,3 +149,25 @@ impl Cheatcode for resumeTracingCall {
         Ok(Default::default())
     }
 }
+impl Cheatcode for setArbitraryStorageCall {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { target } = self;
+        ccx.state.arbitrary_storage.push(*target);
+
+        Ok(Default::default())
+    }
+}
+
+impl Cheatcode for copyStorageCall {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { from, to } = self;
+        if let Ok(from_account) = ccx.load_account(*from) {
+            let from_storage = from_account.storage.clone();
+            if let Ok(mut to_account) = ccx.load_account(*to) {
+                to_account.storage = from_storage;
+            }
+        }
+
+        Ok(Default::default())
+    }
+}

crates/evm/evm/src/inspectors/stack.rs

@@ -726,6 +726,18 @@ impl<'a, DB: DatabaseExt> Inspector<DB> for InspectorStackRefMut<'a> {
 
         ecx.journaled_state.depth += self.in_inner_context as usize;
         if let Some(cheatcodes) = self.cheatcodes.as_deref_mut() {
+            // Handle mocked functions, replace bytecode address with mock if matched.
+            if let Some(mocks) = cheatcodes.mocked_functions.get(&call.target_address) {
+                if let Some(target) = mocks.get(&call.input) {
+                    call.bytecode_address = *target;
+                } else {
+                    // Check if we have a catch-all mock set for selector.
+                    if let Some(target) = mocks.get(&call.input.slice(..4)) {
+                        call.bytecode_address = *target;
+                    }
+                }
+            }
+
             if let Some(output) = cheatcodes.call_with_executor(ecx, call, self.inner) {
                 if output.result.result != InstructionResult::Continue {
                     ecx.journaled_state.depth -= self.in_inner_context as usize;

crates/forge/tests/it/cheats.rs

@@ -7,6 +7,7 @@ use crate::{
         TEST_DATA_MULTI_VERSION,
     },
 };
+use alloy_primitives::U256;
 use foundry_config::{fs_permissions::PathPermission, FsPermissions};
 use foundry_test_utils::Filter;
 
@@ -26,6 +27,7 @@ async fn test_cheats_local(test_data: &ForgeTestData) {
     }
 
     let mut config = test_data.config.clone();
+    config.fuzz.seed = Some(U256::from(100));
     config.fs_permissions = FsPermissions::new(vec![PathPermission::read_write("./")]);
     let runner = test_data.runner_with_config(config);