Add depth trace tag when running forge test

Author: tskoyo

crates/evm/traces/src/lib.rs

@@ -187,6 +187,15 @@ pub fn render_trace_arena(arena: &SparsedTraceArena) -> String {
     render_trace_arena_inner(arena, false, false)
 }
 
+/// Prunes trace depth if depth is provided as an argument
+pub fn prune_trace_depth(arena: &mut CallTraceArena, depth: usize) {
+    for node in arena.nodes_mut() {
+        if node.trace.depth >= depth {
+            Vec::clear(&mut node.ordering);
+        }
+    }
+}
+
 /// Render a collection of call traces to a string optionally including contract creation bytecodes
 /// and in JSON format.
 pub fn render_trace_arena_inner(

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

@@ -41,7 +41,7 @@ use foundry_config::{
 use foundry_debugger::Debugger;
 use foundry_evm::{
     opts::EvmOpts,
-    traces::{backtrace::BacktraceBuilder, identifier::TraceIdentifiers},
+    traces::{backtrace::BacktraceBuilder, identifier::TraceIdentifiers, prune_trace_depth},
 };
 use regex::Regex;
 use std::{
@@ -136,6 +136,10 @@ pub struct TestArgs {
     #[arg(long, short, env = "FORGE_SUPPRESS_SUCCESSFUL_TRACES", help_heading = "Display options")]
     suppress_successful_traces: bool,
 
+    /// Defines the depth of a trace
+    #[arg(long, short)]
+    depth: Option<usize>,
+
     /// Output test results as JUnit XML report.
     #[arg(long, conflicts_with_all = ["quiet", "json", "gas_report", "summary", "list", "show_progress"], help_heading = "Display options")]
     pub junit: bool,
@@ -652,6 +656,11 @@ impl TestArgs {
 
                     if should_include {
                         decode_trace_arena(arena, &decoder).await;
+
+                        if let Some(depth) = self.depth {
+                            prune_trace_depth(arena, depth);
+                        }
+
                         decoded_traces.push(render_trace_arena_inner(arena, false, verbosity > 4));
                     }
                 }
@@ -1037,6 +1046,12 @@ mod tests {
         assert!(args.fuzz_seed.is_some());
     }
 
+    #[test]
+    fn depth_trace() {
+        let args: TestArgs = TestArgs::parse_from(["foundry-cli", "--depth", "2"]);
+        assert!(args.depth.is_some());
+    }
+
     // <https://github.com/foundry-rs/foundry/issues/5913>
     #[test]
     fn fuzz_seed_exists() {