feat: identify internal function invocations in traces

crates/evm/traces/src/debug/mod.rs

@@ -258,7 +258,7 @@ fn try_decode_args_from_step(args: &Parameters<'_>, step: &CallTraceStep) -> Opt
                         (DynSolType::Uint(8), Some(Storage::Memory | Storage::Storage)) => None,
                         (_, Some(Storage::Memory)) => decode_from_memory(
                             type_,
-                            step.memory.as_ref().unwrap().as_bytes(),
+                            step.memory.as_ref()?.as_bytes(),
                             input.try_into().ok()?,
                         ),
                         // Read other types from stack

crates/evm/traces/src/lib.rs

@@ -118,7 +118,11 @@ pub enum TraceMode {
     Call,
     /// Call trace with tracing for JUMP and JUMPDEST opcode steps.
     ///
-    /// Used for internal functions identification.
+    /// Used for internal functions identification. Does not track memory snapshots.
+    JumpSimple,
+    /// Call trace with tracing for JUMP and JUMPDEST opcode steps.
+    ///
+    /// Same as `JumpSimple`, but tracks memory snapshots as well.
     Jump,
     /// Call trace with complete steps tracing.
     ///
@@ -135,6 +139,10 @@ impl TraceMode {
         matches!(self, Self::Call)
     }
 
+    pub const fn is_jump_simple(self) -> bool {
+        matches!(self, Self::JumpSimple)
+    }
+
     pub const fn is_jump(self) -> bool {
         matches!(self, Self::Jump)
     }
@@ -159,6 +167,14 @@ impl TraceMode {
         }
     }
 
+    pub fn with_decode_internal_simple(self, yes: bool) -> Self {
+        if yes {
+            std::cmp::max(self, Self::JumpSimple)
+        } else {
+            self
+        }
+    }
+
     pub fn with_verbosity(self, verbosiy: u8) -> Self {
         if verbosiy >= 3 {
             std::cmp::max(self, Self::Call)
@@ -172,18 +188,17 @@ impl TraceMode {
             None
         } else {
             TracingInspectorConfig {
-                record_steps: self.is_debug() || self.is_jump(),
-                record_memory_snapshots: self.is_debug() || self.is_jump(),
-                record_stack_snapshots: if self.is_debug() || self.is_jump() {
+                record_steps: self >= Self::JumpSimple,
+                record_memory_snapshots: self >= Self::Jump,
+                record_stack_snapshots: if self >= Self::JumpSimple {
                     StackSnapshotType::Full
                 } else {
                     StackSnapshotType::None
                 },
                 record_logs: true,
                 record_state_diff: false,
                 record_returndata_snapshots: self.is_debug(),
-                record_opcodes_filter: self
-                    .is_jump()
+                record_opcodes_filter: (self.is_jump() || self.is_jump_simple())
                     .then(|| OpcodeFilter::new().enabled(OpCode::JUMP).enabled(OpCode::JUMPDEST)),
                 exclude_precompile_calls: false,
             }

crates/forge/bin/cmd/test/mod.rs

@@ -84,7 +84,7 @@ pub struct TestArgs {
     /// If more than one test matches your specified criteria, you must add additional filters
     /// until only one test is found (see --match-contract and --match-path).
     #[arg(long, value_name = "TEST_FUNCTION")]
-    decode_internal: Option<Regex>,
+    decode_internal: Option<Option<Regex>>,
 
     /// Print a gas report.
     #[arg(long, env = "FORGE_GAS_REPORT")]
@@ -307,12 +307,20 @@ impl TestArgs {
 
         let env = evm_opts.evm_env().await?;
 
+        // If we are provided with test function regex, we are enabling complete internal fns
+        // tracing.
+        let decode_internal = self.decode_internal.as_ref().map_or(false, |v| v.is_some());
+        // If we are provided with just --decode-internal flag, we enable simple tracing (without
+        // memory decoding).
+        let decode_internal_simple = self.decode_internal.is_some();
+
         // Prepare the test builder.
         let should_debug = self.debug.is_some();
         let config = Arc::new(config);
         let runner = MultiContractRunnerBuilder::new(config.clone())
             .set_debug(should_debug)
-            .set_decode_internal(self.decode_internal.is_some())
+            .set_decode_internal(decode_internal)
+            .set_decode_internal_simple(decode_internal_simple)
             .initial_balance(evm_opts.initial_balance)
             .evm_spec(config.evm_spec_id())
             .sender(evm_opts.sender)
@@ -337,7 +345,10 @@ impl TestArgs {
         };
 
         maybe_override_mt("debug", self.debug.as_ref())?;
-        maybe_override_mt("decode_internal", self.decode_internal.as_ref())?;
+        maybe_override_mt(
+            "decode-internal",
+            self.decode_internal.as_ref().and_then(|v| v.as_ref()),
+        )?;
 
         let libraries = runner.libraries.clone();
         let outcome = self.run_tests(runner, config, verbosity, &filter, &output).await?;
@@ -391,7 +402,9 @@ impl TestArgs {
         trace!(target: "forge::test", "running all tests");
 
         let num_filtered = runner.matching_test_functions(filter).count();
-        if (self.debug.is_some() || self.decode_internal.is_some()) && num_filtered != 1 {
+        if (self.debug.is_some() || self.decode_internal.as_ref().map_or(false, |v| v.is_some())) &&
+            num_filtered != 1
+        {
             eyre::bail!(
                 "{num_filtered} tests matched your criteria, but exactly 1 test must match in order to run the debugger.\n\n\
                  Use --match-contract and --match-path to further limit the search.\n\

crates/forge/src/multi_runner.rs

@@ -62,6 +62,8 @@ pub struct MultiContractRunner {
     pub debug: bool,
     /// Whether to enable steps tracking in the tracer.
     pub decode_internal: bool,
+    /// Whether to enable simple steps tracing (without memory recording).
+    pub decode_internal_simple: bool,
     /// Settings related to fuzz and/or invariant tests
     pub test_options: TestOptions,
     /// Whether to enable call isolation
@@ -240,6 +242,7 @@ impl MultiContractRunner {
 
         let trace_mode = TraceMode::default()
             .with_debug(self.debug)
+            .with_decode_internal_simple(self.decode_internal_simple)
             .with_decode_internal(self.decode_internal)
             .with_verbosity(self.evm_opts.verbosity);
 
@@ -307,6 +310,8 @@ pub struct MultiContractRunnerBuilder {
     pub debug: bool,
     /// Whether to enable steps tracking in the tracer.
     pub decode_internal: bool,
+    /// Whether to enable simple steps tracing (without memory recording).
+    pub decode_internal_simple: bool,
     /// Whether to enable call isolation
     pub isolation: bool,
     /// Settings related to fuzz and/or invariant tests
@@ -326,6 +331,7 @@ impl MultiContractRunnerBuilder {
             isolation: Default::default(),
             test_options: Default::default(),
             decode_internal: Default::default(),
+            decode_internal_simple: Default::default(),
         }
     }
 
@@ -369,6 +375,11 @@ impl MultiContractRunnerBuilder {
         self
     }
 
+    pub fn set_decode_internal_simple(mut self, enable: bool) -> Self {
+        self.decode_internal_simple = enable;
+        self
+    }
+
     pub fn enable_isolation(mut self, enable: bool) -> Self {
         self.isolation = enable;
         self
@@ -440,6 +451,7 @@ impl MultiContractRunnerBuilder {
             coverage: self.coverage,
             debug: self.debug,
             decode_internal: self.decode_internal,
+            decode_internal_simple: self.decode_internal_simple,
             test_options: self.test_options.unwrap_or_default(),
             isolation: self.isolation,
             known_contracts,

crates/forge/tests/cli/test_cmd.rs

@@ -3,7 +3,7 @@
 use alloy_primitives::U256;
 use foundry_config::{Config, FuzzConfig};
 use foundry_test_utils::{
-    rpc,
+    rpc, str,
     util::{OutputExt, OTHER_SOLC_VERSION, SOLC_VERSION},
 };
 use std::{path::PathBuf, str::FromStr};
@@ -853,3 +853,120 @@ forgetest_init!(should_not_show_logs_when_fuzz_test_inline_config, |prj, cmd| {
     let stdout = cmd.stdout_lossy();
     assert!(!stdout.contains("inside fuzz test, x is:"), "\n{stdout}");
 });
+
+// tests internal functions trace
+forgetest_init!(internal_functions_trace, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.clear();
+
+    // Disable optimizer because for simple contract most functions will get inlined.
+    prj.write_config(Config { optimizer: false, ..Default::default() });
+
+    prj.add_test(
+        "Simple",
+        r#"pragma solidity 0.8.24;
+        import {Test, console2} from "forge-std/Test.sol";
+contract SimpleContract {
+    uint256 public num;
+    address public addr;
+
+    function _setNum(uint256 _num) internal returns(uint256 prev) {
+        prev = num;
+        num = _num;
+    }
+
+    function _setAddr(address _addr) internal returns(address prev) {
+        prev = addr;
+        addr = _addr;
+    }
+
+    function increment() public {
+        _setNum(num + 1);
+    }
+
+    function setValues(uint256 _num, address _addr) public {
+        _setNum(_num);
+        _setAddr(_addr);
+    }
+}
+
+contract SimpleContractTest is Test {
+    function test() public {
+        SimpleContract c = new SimpleContract();
+        c.increment();
+        c.setValues(100, address(0x123));
+    }
+}
+     "#,
+    )
+    .unwrap();
+    cmd.args(["test", "-vvvv", "--decode-internal"]).assert_success().stdout_eq(str![[r#"
+...
+Traces:
+  [250463] SimpleContractTest::test()
+    ├─ [171014] → new SimpleContract@0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f
+    │   └─ ← [Return] 854 bytes of code
+    ├─ [22638] SimpleContract::increment()
+    │   ├─ [20150] SimpleContract::_setNum(1)
+    │   │   └─ ← 0
+    │   └─ ← [Stop] 
+    ├─ [23219] SimpleContract::setValues(100, 0x0000000000000000000000000000000000000123)
+    │   ├─ [250] SimpleContract::_setNum(100)
+    │   │   └─ ← 1
+    │   ├─ [22339] SimpleContract::_setAddr(0x0000000000000000000000000000000000000123)
+    │   │   └─ ← 0x0000000000000000000000000000000000000000
+    │   └─ ← [Stop] 
+    └─ ← [Stop] 
+...
+"#]]);
+});
+
+// tests internal functions trace with memory decoding
+forgetest_init!(internal_functions_trace_memory, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.clear();
+
+    // Disable optimizer because for simple contract most functions will get inlined.
+    prj.write_config(Config { optimizer: false, ..Default::default() });
+
+    prj.add_test(
+        "Simple",
+        r#"pragma solidity 0.8.24;
+import {Test, console2} from "forge-std/Test.sol";
+
+contract SimpleContract {
+    string public str = "initial value";
+
+    function _setStr(string memory _str) internal returns(string memory prev) {
+        prev = str;
+        str = _str;
+    }
+
+    function setStr(string memory _str) public {
+        _setStr(_str);
+    }
+}
+
+contract SimpleContractTest is Test {
+    function test() public {
+        SimpleContract c = new SimpleContract();
+        c.setStr("new value");
+    }
+}
+     "#,
+    )
+    .unwrap();
+    cmd.args(["test", "-vvvv", "--decode-internal", "test"]).assert_success().stdout_eq(str![[r#"
+...
+Traces:
+  [421960] SimpleContractTest::test()
+    ├─ [385978] → new SimpleContract@0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f
+    │   └─ ← [Return] 1814 bytes of code
+    ├─ [2534] SimpleContract::setStr("new value")
+    │   ├─ [1600] SimpleContract::_setStr("new value")
+    │   │   └─ ← "initial value"
+    │   └─ ← [Stop] 
+    └─ ← [Stop] 
+...
+"#]]);
+});