diff --git a/crates/evm/traces/src/lib.rs b/crates/evm/traces/src/lib.rs index f09b8d756b93b..301454d43eeab 100644 --- a/crates/evm/traces/src/lib.rs +++ b/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 { + node.ordering.clear(); + } + } +} + /// Render a collection of call traces to a string optionally including contract creation bytecodes /// and in JSON format. pub fn render_trace_arena_inner( diff --git a/crates/forge/src/cmd/test/mod.rs b/crates/forge/src/cmd/test/mod.rs index 0a45e1574f1bd..a37a422e6688b 100644 --- a/crates/forge/src/cmd/test/mod.rs +++ b/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)] + trace_depth: Option, + /// 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(trace_depth) = self.trace_depth { + prune_trace_depth(arena, trace_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", "--trace-depth", "2"]); + assert!(args.trace_depth.is_some()); + } + // #[test] fn fuzz_seed_exists() {