feat(cheatcode): startDebugTraceRecording and stopDebugTraceRecording for ERC4337 testing

crates/cheatcodes/Cargo.toml

@@ -58,6 +58,7 @@ p256 = "0.13.2"
 ecdsa = "0.16"
 rand = "0.8"
 revm.workspace = true
+revm-inspectors.workspace = true
 semver.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true

crates/cheatcodes/spec/src/lib.rs

@@ -85,6 +85,7 @@ impl Cheatcodes<'static> {
                 Vm::AccountAccess::STRUCT.clone(),
                 Vm::StorageAccess::STRUCT.clone(),
                 Vm::Gas::STRUCT.clone(),
+                Vm::DebugStep::STRUCT.clone(),
             ]),
             enums: Cow::Owned(vec![
                 Vm::CallerMode::ENUM.clone(),

crates/cheatcodes/spec/src/vm.rs

@@ -261,6 +261,28 @@ interface Vm {
         uint64 depth;
     }
 
+    /// The result of the `stopDebugTraceRecording` call
+    struct DebugStep {
+        /// The stack before executing the step of the run.
+        /// stack\[0\] represents the top of the stack.
+        /// and only stack data relevant to the opcode execution is contained.
+        uint256[] stack;
+        /// The memory input data before executing the step of the run.
+        /// only input data relevant to the opcode execution is contained.
+        ///
+        /// e.g. for MLOAD, it will have memory\[offset:offset+32\] copied here.
+        /// the offset value can be get by the stack data.
+        bytes memoryInput;
+        /// The opcode that was accessed.
+        uint8 opcode;
+        /// The call depth of the step.
+        uint64 depth;
+        /// Whether the call end up with out of gas error.
+        bool isOutOfGas;
+        /// The contract address where the opcode is running
+        address contractAddr;
+    }
+
     // ======== EVM ========
 
     /// Gets the address for a given private key.
@@ -287,6 +309,17 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Unsafe)]
     function loadAllocs(string calldata pathToAllocsJson) external;
 
+    // -------- Record Debug Traces --------
+
+    /// Records the debug trace during the run.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function startDebugTraceRecording() external;
+
+    /// Stop debug trace recording and returns the recorded debug trace.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function stopAndReturnDebugTraceRecording() external returns (DebugStep[] memory step);
+
+
     /// Clones a source account code, state, balance and nonce to a target account and updates in-memory EVM state.
     #[cheatcode(group = Evm, safety = Unsafe)]
     function cloneAccount(address source, address target) external;

crates/cheatcodes/src/evm.rs

@@ -1,8 +1,9 @@
 //! Implementations of [`Evm`](spec::Group::Evm) cheatcodes.
 
 use crate::{
-    inspector::InnerEcx, BroadcastableTransaction, Cheatcode, Cheatcodes, CheatcodesExecutor,
-    CheatsCtxt, Result, Vm::*,
+    inspector::{InnerEcx, RecordDebugStepInfo},
+    BroadcastableTransaction, Cheatcode, Cheatcodes, CheatcodesExecutor, CheatsCtxt, Error, Result,
+    Vm::*,
 };
 use alloy_consensus::TxEnvelope;
 use alloy_genesis::{Genesis, GenesisAccount};
@@ -14,10 +15,14 @@ use foundry_evm_core::{
     backend::{DatabaseExt, RevertStateSnapshotAction},
     constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS},
 };
+use foundry_evm_traces::StackSnapshotType;
 use rand::Rng;
 use revm::primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY};
 use std::{collections::BTreeMap, path::Path};
 
+mod record_debug_step;
+use record_debug_step::{convert_call_trace_to_debug_step, flatten_call_trace};
+
 mod fork;
 pub(crate) mod mapping;
 pub(crate) mod mock;
@@ -715,6 +720,64 @@ impl Cheatcode for setBlockhashCall {
     }
 }
 
+impl Cheatcode for startDebugTraceRecordingCall {
+    fn apply_full(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result {
+        let Some(tracer) = executor.tracing_inspector().and_then(|t| t.as_mut()) else {
+            return Err(Error::from("no tracer initiated, consider adding -vvv flag"))
+        };
+
+        let mut info = RecordDebugStepInfo {
+            // will be updated later
+            start_node_idx: 0,
+            // keep the original config to revert back later
+            original_tracer_config: *tracer.config(),
+        };
+
+        // turn on tracer configuration for recording
+        tracer.update_config(|config| {
+            config
+                .set_steps(true)
+                .set_memory_snapshots(true)
+                .set_stack_snapshots(StackSnapshotType::Full)
+        });
+
+        // track where the recording starts
+        if let Some(last_node) = tracer.traces().nodes().last() {
+            info.start_node_idx = last_node.idx;
+        }
+
+        ccx.state.record_debug_steps_info = Some(info);
+        Ok(Default::default())
+    }
+}
+
+impl Cheatcode for stopAndReturnDebugTraceRecordingCall {
+    fn apply_full(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result {
+        let Some(tracer) = executor.tracing_inspector().and_then(|t| t.as_mut()) else {
+            return Err(Error::from("no tracer initiated, consider adding -vvv flag"))
+        };
+
+        let Some(record_info) = ccx.state.record_debug_steps_info else {
+            return Err(Error::from("nothing recorded"))
+        };
+
+        // Revert the tracer config to the one before recording
+        tracer.update_config(|_config| record_info.original_tracer_config);
+
+        // Use the trace nodes to flatten the call trace
+        let root = tracer.traces();
+        let steps = flatten_call_trace(0, root, record_info.start_node_idx);
+
+        let debug_steps: Vec<DebugStep> =
+            steps.iter().map(|&step| convert_call_trace_to_debug_step(step)).collect();
+
+        // Clean up the recording info
+        ccx.state.record_debug_steps_info = None;
+
+        Ok(debug_steps.abi_encode())
+    }
+}
+
 pub(super) fn get_nonce(ccx: &mut CheatsCtxt, address: &Address) -> Result {
     let account = ccx.ecx.journaled_state.load_account(*address, &mut ccx.ecx.db)?;
     Ok(account.info.nonce.abi_encode())