diff --git a/Cargo.lock b/Cargo.lock index cc355c59bc22..93a4a610d74d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2891,7 +2891,7 @@ dependencies = [ "console_error_panic_hook", "pest", "pest_derive", - "quick-xml", + "quick-xml 0.18.1", "wasm-bindgen", ] @@ -2946,6 +2946,15 @@ dependencies = [ "regex", ] +[[package]] +name = "env_logger" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd405aab171cb85d6735e5c8d9db038c17d3ca007a4d2c25f337935c3d90580" +dependencies = [ + "log", +] + [[package]] name = "env_logger" version = "0.11.3" @@ -3341,6 +3350,7 @@ dependencies = [ "humantime-serde", "hyper 1.3.1", "indicatif", + "inferno", "itertools 0.13.0", "mockall", "once_cell", @@ -5040,6 +5050,29 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "inferno" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "321f0f839cd44a4686e9504b0a62b4d69a50b62072144c71c68f5873c167b8d9" +dependencies = [ + "ahash", + "clap", + "crossbeam-channel", + "crossbeam-utils", + "dashmap", + "env_logger 0.10.2", + "indexmap 2.2.6", + "is-terminal", + "itoa", + "log", + "num-format", + "once_cell", + "quick-xml 0.26.0", + "rgb", + "str_stack", +] + [[package]] name = "inlinable_string" version = "0.1.15" @@ -5457,7 +5490,7 @@ dependencies = [ "clap", "clap_complete", "elasticlunr-rs", - "env_logger", + "env_logger 0.11.3", "handlebars", "log", "memchr", @@ -6646,6 +6679,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "quick-xml" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f50b1c63b38611e7d4d7f68b82d3ad0cc71a2ad2e7f61fc10f1328d917c93cd" +dependencies = [ + "memchr", +] + [[package]] name = "quinn" version = "0.11.2" @@ -7068,6 +7110,15 @@ dependencies = [ "subtle", ] +[[package]] +name = "rgb" +version = "0.8.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7439be6844e40133eda024efd85bf07f59d0dd2f59b10c00dd6cfb92cc5c741" +dependencies = [ + "bytemuck", +] + [[package]] name = "ring" version = "0.17.8" @@ -8016,6 +8067,12 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e08d8363704e6c71fc928674353e6b7c23dcea9d82d7012c8faf2a3a025f8d0" +[[package]] +name = "str_stack" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9091b6114800a5f2141aee1d1b9d6ca3592ac062dc5decb3764ec5895a47b4eb" + [[package]] name = "string_cache" version = "0.8.7" diff --git a/crates/chisel/src/dispatcher.rs b/crates/chisel/src/dispatcher.rs index 5c13c8e01f8d..f688b5e191e1 100644 --- a/crates/chisel/src/dispatcher.rs +++ b/crates/chisel/src/dispatcher.rs @@ -935,7 +935,7 @@ impl ChiselDispatcher { for (kind, trace) in &result.traces { // Display all Setup + Execution traces. if matches!(kind, TraceKind::Setup | TraceKind::Execution) { - println!("{}", render_trace_arena(trace, decoder).await?); + println!("{}", render_trace_arena(trace, decoder).await?.0); } } diff --git a/crates/cli/src/utils/cmd.rs b/crates/cli/src/utils/cmd.rs index ac5a5ee610a9..42c78d4518e1 100644 --- a/crates/cli/src/utils/cmd.rs +++ b/crates/cli/src/utils/cmd.rs @@ -412,7 +412,7 @@ pub async fn print_traces(result: &mut TraceResult, decoder: &CallTraceDecoder) println!("Traces:"); for (_, arena) in traces { - println!("{}", render_trace_arena(arena, decoder).await?); + println!("{}", render_trace_arena(arena, decoder).await?.0); } println!(); diff --git a/crates/evm/traces/src/folded_stack_trace.rs b/crates/evm/traces/src/folded_stack_trace.rs new file mode 100644 index 000000000000..b9b6c63bc272 --- /dev/null +++ b/crates/evm/traces/src/folded_stack_trace.rs @@ -0,0 +1,150 @@ +#[derive(Debug, Clone, Default)] +pub struct FoldedStackTrace { + traces: Vec<(Vec, i64)>, + exits: Option, +} + +impl FoldedStackTrace { + pub fn enter(&mut self, label: String, gas: i64) { + let mut trace_entry = self.traces.last().map(|entry| entry.0.clone()).unwrap_or_default(); + + let mut exits = self.exits.unwrap_or_default(); + while exits > 0 { + trace_entry.pop(); + exits -= 1; + } + self.exits = None; + + trace_entry.push(label); + self.traces.push((trace_entry, gas)); + } + + pub fn exit(&mut self) { + self.exits = self.exits.map(|exits| exits + 1).or(Some(1)); + } + + pub fn fold(mut self) -> Vec { + self.subtract_children(); + self.fold_without_subtraction() + } + + pub fn fold_without_subtraction(&mut self) -> Vec { + let mut lines = Vec::new(); + for (trace, gas) in self.traces.iter() { + lines.push(format!("{} {}", trace.join(";"), gas)); + } + lines + } + + pub fn subtract_children(&mut self) { + // Iterate over each trace to find the children and subtract their values from the parents + for i in 0..self.traces.len() { + let (left, right) = self.traces.split_at_mut(i); + let (trace, gas) = &right[0]; + if trace.len() > 1 { + let parent_trace_to_match = &trace[..trace.len() - 1]; + for parent in left { + if parent.0 == parent_trace_to_match { + parent.1 -= gas; + break; + } + } + } + } + } +} + +mod tests { + #[test] + fn test_insert_1() { + let mut trace = super::FoldedStackTrace::default(); + trace.enter("top".to_string(), 500); + trace.enter("child_a".to_string(), 100); + trace.exit(); + trace.enter("child_b".to_string(), 200); + + assert_eq!( + trace.fold_without_subtraction(), + vec![ + "top 500", // + "top;child_a 100", + "top;child_b 200", + ] + ); + assert_eq!( + trace.fold(), + vec![ + "top 200", // 500 - 100 - 200 + "top;child_a 100", + "top;child_b 200", + ] + ); + } + + #[test] + fn test_insert_2() { + let mut trace = super::FoldedStackTrace::default(); + trace.enter("top".to_string(), 500); + trace.enter("child_a".to_string(), 300); + trace.enter("child_b".to_string(), 100); + trace.exit(); + trace.exit(); + trace.enter("child_c".to_string(), 100); + + assert_eq!( + trace.fold_without_subtraction(), + vec![ + "top 500", // + "top;child_a 300", + "top;child_a;child_b 100", + "top;child_c 100", + ] + ); + + assert_eq!( + trace.fold(), + vec![ + "top 100", // 500 - 300 - 100 + "top;child_a 200", // 300 - 100 + "top;child_a;child_b 100", + "top;child_c 100", + ] + ); + } + + #[test] + fn test_insert_3() { + let mut trace = super::FoldedStackTrace::default(); + trace.enter("top".to_string(), 1700); + trace.enter("child_a".to_string(), 500); + trace.exit(); + trace.enter("child_b".to_string(), 500); + trace.enter("child_c".to_string(), 500); + trace.exit(); + trace.exit(); + trace.exit(); + trace.enter("top2".to_string(), 1700); + + assert_eq!( + trace.fold_without_subtraction(), + vec![ + "top 1700", // + "top;child_a 500", + "top;child_b 500", + "top;child_b;child_c 500", + "top2 1700", + ] + ); + + assert_eq!( + trace.fold(), + vec![ + "top 700", // + "top;child_a 500", + "top;child_b 0", + "top;child_b;child_c 500", + "top2 1700", + ] + ); + } +} diff --git a/crates/evm/traces/src/lib.rs b/crates/evm/traces/src/lib.rs index f70ab78abd91..430ff4aeba57 100644 --- a/crates/evm/traces/src/lib.rs +++ b/crates/evm/traces/src/lib.rs @@ -35,6 +35,9 @@ pub use decoder::{CallTraceDecoder, CallTraceDecoderBuilder}; pub mod debug; +pub mod folded_stack_trace; +use folded_stack_trace::FoldedStackTrace; + pub type Traces = Vec<(TraceKind, CallTraceArena)>; #[derive(Default, Debug, Eq, PartialEq)] @@ -81,10 +84,10 @@ const RETURN: &str = "← "; /// Render a collection of call traces. /// /// The traces will be decoded using the given decoder, if possible. -pub async fn render_trace_arena( +pub async fn render_trace_arena<'a>( arena: &CallTraceArena, decoder: &CallTraceDecoder, -) -> Result { +) -> Result<(String, Vec), std::fmt::Error> { decoder.prefetch_signatures(arena.nodes()).await; let identified_internals = &decoder.identify_arena_steps(arena); @@ -95,6 +98,7 @@ pub async fn render_trace_arena( decoder: &'a CallTraceDecoder, identified_internals: &'a [Vec], s: &'a mut String, + folded_stack_traces: &'a mut FoldedStackTrace, node_idx: usize, mut ordering_idx: usize, internal_end_step_idx: Option, @@ -126,6 +130,7 @@ pub async fn render_trace_arena( decoder, &identified_internals, s, + folded_stack_traces, node.children[*index], left, right, @@ -153,6 +158,8 @@ pub async fn render_trace_arena( .map(|v| format!("({})", v.join(", "))) .unwrap_or_default() )?; + folded_stack_traces + .enter(decoded.function_name.to_string(), decoded.gas_used); let left_prefix = format!("{right}{BRANCH}"); let right_prefix = format!("{right}{PIPE}"); ordering_idx = render_items( @@ -160,6 +167,7 @@ pub async fn render_trace_arena( decoder, identified_internals, s, + folded_stack_traces, node_idx, ordering_idx + 1, decoded.end_step_idx, @@ -175,6 +183,7 @@ pub async fn render_trace_arena( } writeln!(s)?; + folded_stack_traces.exit(); } } } @@ -186,11 +195,13 @@ pub async fn render_trace_arena( .boxed() } + #[allow(clippy::too_many_arguments)] fn inner<'a>( arena: &'a [CallTraceNode], decoder: &'a CallTraceDecoder, identified_internals: &'a [Vec], s: &'a mut String, + folded_stack_traces: &'a mut FoldedStackTrace, idx: usize, left: &'a str, child: &'a str, @@ -199,7 +210,8 @@ pub async fn render_trace_arena( let node = &arena[idx]; // Display trace header - let (trace, return_data) = render_trace(&node.trace, decoder).await?; + let (trace, return_data) = + render_trace(&node.trace, decoder, folded_stack_traces).await?; writeln!(s, "{left}{trace}")?; // Display logs and subcalls @@ -208,6 +220,7 @@ pub async fn render_trace_arena( decoder, identified_internals, s, + folded_stack_traces, idx, 0, None, @@ -234,14 +247,28 @@ pub async fn render_trace_arena( }?; writeln!(s)?; + folded_stack_traces.exit(); Ok(()) } .boxed() } let mut s = String::new(); - inner(arena.nodes(), decoder, identified_internals, &mut s, 0, " ", " ").await?; - Ok(s) + let mut folded_stack_traces = FoldedStackTrace::default(); + inner( + arena.nodes(), + decoder, + identified_internals, + &mut s, + &mut folded_stack_traces, + 0, + " ", + " ", + ) + .await?; + + let folded = folded_stack_traces.fold(); + Ok((s, folded)) } /// Render a call trace. @@ -250,6 +277,7 @@ pub async fn render_trace_arena( pub async fn render_trace( trace: &CallTrace, decoder: &CallTraceDecoder, + folded_stack_traces: &mut FoldedStackTrace, ) -> Result<(String, Option), std::fmt::Error> { let mut s = String::new(); write!(&mut s, "[{}] ", trace.gas_used)?; @@ -265,6 +293,10 @@ pub async fn render_trace( decoded.label.as_deref().unwrap_or(""), address )?; + folded_stack_traces.enter( + decoded.label.as_deref().unwrap_or("").to_string(), + trace.gas_used as i64, + ); } else { let (func_name, inputs) = match &decoded.func { Some(DecodedCallData { signature, args }) => { @@ -304,6 +336,10 @@ pub async fn render_trace( }, action = action.yellow(), )?; + folded_stack_traces.enter( + format!("{addr}::{func_name}", addr = decoded.label.as_deref().unwrap_or(&address),), + trace.gas_used as i64, + ); } Ok((s, decoded.return_data)) diff --git a/crates/forge/Cargo.toml b/crates/forge/Cargo.toml index 93b15ed61d00..623e9d54d6c7 100644 --- a/crates/forge/Cargo.toml +++ b/crates/forge/Cargo.toml @@ -79,6 +79,7 @@ dunce.workspace = true futures.workspace = true indicatif = "0.17" itertools.workspace = true +inferno = "0.11.19" once_cell.workspace = true parking_lot.workspace = true regex = { version = "1", default-features = false } diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index c30a73286c8b..9eb6fb7a3d29 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -10,7 +10,7 @@ use forge::{ traces::{ debug::{ContractSources, DebugTraceIdentifier}, identifier::SignaturesIdentifier, - CallTraceDecoderBuilder, TraceKind, + CallTraceDecoderBuilder, CallTraceNode, DecodedCallTrace, TraceKind, }, MultiContractRunner, MultiContractRunnerBuilder, TestFilter, TestOptions, TestOptionsBuilder, }; @@ -38,6 +38,7 @@ use foundry_evm::traces::identifier::TraceIdentifiers; use regex::Regex; use std::{ collections::{BTreeMap, BTreeSet}, + io::Read, path::PathBuf, sync::{mpsc::channel, Arc}, time::Instant, @@ -77,6 +78,10 @@ pub struct TestArgs { #[arg(long, value_name = "TEST_FUNCTION")] debug: Option, + /// Generate flamegraph for a test function. Works similar to debug flag. + #[arg(long, value_name = "TEST_FUNCTION")] + flamegraph: Option, + /// Whether to identify internal functions in traces. #[arg(long)] decode_internal: bool, @@ -289,16 +294,22 @@ impl TestArgs { .profiles(profiles) .build(&output, project_root)?; + // Prepare the test builder + let should_debug = self.debug.is_some(); + let should_flamegraph = self.flamegraph.is_some(); + // Determine print verbosity and executor verbosity let verbosity = evm_opts.verbosity; - if self.gas_report && evm_opts.verbosity < 3 { + if (self.gas_report && evm_opts.verbosity < 3) || should_flamegraph { evm_opts.verbosity = 3; } let env = evm_opts.evm_env().await?; - // Prepare the test builder - let should_debug = self.debug.is_some(); + // Clone the output only if we actually need it later for the debugger. + let output_clone = + (should_debug || should_flamegraph || self.decode_internal).then(|| output.clone()); + let config = Arc::new(config); let runner = MultiContractRunnerBuilder::new(config.clone()) .set_debug(should_debug) @@ -311,20 +322,81 @@ impl TestArgs { .enable_isolation(evm_opts.isolate) .build(project_root, &output, env, evm_opts)?; - if let Some(debug_test_pattern) = &self.debug { + if should_debug && should_flamegraph { + eyre::bail!("Cannot specify both --debug and --flamegraph."); + } + + let mut set_test_pattern = |new_test_pattern: &Regex, flag_name| { let test_pattern = &mut filter.args_mut().test_pattern; if test_pattern.is_some() { eyre::bail!( - "Cannot specify both --debug and --match-test. \ - Use --match-contract and --match-path to further limit the search instead." + "Cannot specify both --{} and --match-test. \ + Use --match-contract and --match-path to further limit the search instead.", + flag_name ); } - *test_pattern = Some(debug_test_pattern.clone()); + *test_pattern = Some(new_test_pattern.clone()); + Ok(()) + }; + + if let Some(debug_test_pattern) = &self.debug { + set_test_pattern(debug_test_pattern, "debug")?; + } + + if let Some(flamegraph_test_pattern) = &self.flamegraph { + set_test_pattern(flamegraph_test_pattern, "flamegraph")?; } let libraries = runner.libraries.clone(); let outcome = self.run_tests(runner, config, verbosity, &filter, &output).await?; + if should_flamegraph { + let (_, test_result) = outcome + .results + .iter() + .find(|(_, r)| !r.test_results.is_empty()) + .map(|(_, r)| (r, r.test_results.values().next().unwrap())) + .unwrap(); + + let arena = test_result + .traces + .iter() + .find_map( + |(kind, arena)| { + if *kind == TraceKind::Execution { + Some(arena) + } else { + None + } + }, + ) + .unwrap(); + + let decoder = outcome.last_run_decoder.as_ref().unwrap(); + decoder.prefetch_signatures(arena.nodes()).await; + + let mut folded_stack_lines = render_trace_arena(arena, decoder).await?.1; + folded_stack_lines.reverse(); + + let file_name = "flamegraph.svg"; + let writer = std::fs::File::create(file_name).unwrap(); + + let mut options = inferno::flamegraph::Options::default(); + options.flame_chart = true; + inferno::flamegraph::from_lines( + &mut options, + folded_stack_lines.iter().map(|s| s.as_str()), + writer, + ) + .unwrap(); + + let mut buf = String::new(); + let mut file = std::fs::File::open(file_name).unwrap(); + file.read_to_string(&mut buf).expect("failed to read flamegraph file"); + let buf = buf.replace("samples", "gas"); + std::fs::write(file_name, buf).expect("failed to write flamegraph file"); + } + if should_debug { // Get first non-empty suite result. We will have only one such entry let Some((_, test_result)) = outcome @@ -374,11 +446,12 @@ impl TestArgs { trace!(target: "forge::test", "running all tests"); let num_filtered = runner.matching_test_functions(filter).count(); - if self.debug.is_some() && num_filtered != 1 { + if (self.debug.is_some() || self.flamegraph.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\ + "{num_filtered} tests matched your criteria, but exactly 1 test must match in order to {action}.\n\n\ Use --match-contract and --match-path to further limit the search.\n\ - Filter used:\n{filter}" + Filter used:\n{filter}", + action = if self.debug.is_some() {"run the debugger"} else {"generate a flamegraph"}, ); } @@ -444,7 +517,10 @@ impl TestArgs { decoder.clear_addresses(); // We identify addresses if we're going to print *any* trace or gas report. - let identify_addresses = verbosity >= 3 || self.gas_report || self.debug.is_some(); + let identify_addresses = verbosity >= 3 || + self.gas_report || + self.debug.is_some() || + self.flamegraph.is_some(); // Print suite header. println!(); @@ -507,7 +583,7 @@ impl TestArgs { }; if should_include { - decoded_traces.push(render_trace_arena(arena, &decoder).await?); + decoded_traces.push(render_trace_arena(arena, &decoder).await?.0); } } diff --git a/crates/script/src/execute.rs b/crates/script/src/execute.rs index 3552e3977f97..60d6e186cf17 100644 --- a/crates/script/src/execute.rs +++ b/crates/script/src/execute.rs @@ -428,7 +428,7 @@ impl PreSimulationState { } || !result.success; if should_include { - shell::println(render_trace_arena(trace, decoder).await?)?; + shell::println(render_trace_arena(trace, decoder).await?.0)?; } } shell::println(String::new())?; diff --git a/crates/script/src/simulate.rs b/crates/script/src/simulate.rs index d2546cda3899..d26910907346 100644 --- a/crates/script/src/simulate.rs +++ b/crates/script/src/simulate.rs @@ -161,7 +161,7 @@ impl PreSimulationState { for (_, trace) in &traces { println!( "{}", - render_trace_arena(trace, &self.execution_artifacts.decoder).await? + render_trace_arena(trace, &self.execution_artifacts.decoder).await?.0 ); } }