diff --git a/Cargo.lock b/Cargo.lock index c353292abfe05..818e7367943d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2898,7 +2898,7 @@ dependencies = [ "console_error_panic_hook", "pest", "pest_derive", - "quick-xml", + "quick-xml 0.18.1", "wasm-bindgen", ] @@ -3337,6 +3337,7 @@ dependencies = [ "humantime-serde", "hyper 1.4.1", "indicatif", + "inferno", "itertools 0.13.0", "mockall", "opener", @@ -5085,6 +5086,23 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "inferno" +version = "0.11.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "232929e1d75fe899576a3d5c7416ad0d88dbfbb3c3d6aa00873a7408a50ddb88" +dependencies = [ + "ahash", + "is-terminal", + "itoa", + "log", + "num-format", + "once_cell", + "quick-xml 0.26.0", + "rgb", + "str_stack", +] + [[package]] name = "inlinable_string" version = "0.1.15" @@ -6777,6 +6795,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.3" @@ -7149,6 +7176,15 @@ dependencies = [ "subtle", ] +[[package]] +name = "rgb" +version = "0.8.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f86ae463694029097b846d8f99fd5536740602ae00022c0c50c5600720b2f71" +dependencies = [ + "bytemuck", +] + [[package]] name = "ring" version = "0.17.8" @@ -8088,6 +8124,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[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/evm/traces/src/folded_stack_trace.rs b/crates/evm/traces/src/folded_stack_trace.rs new file mode 100644 index 0000000000000..de76f8e804085 --- /dev/null +++ b/crates/evm/traces/src/folded_stack_trace.rs @@ -0,0 +1,301 @@ +use alloy_primitives::hex::ToHexExt; +use revm_inspectors::tracing::{ + types::{CallTraceNode, CallTraceStep, DecodedTraceStep, TraceMemberOrder}, + CallTraceArena, +}; + +/// Builds a folded stack trace from a call trace arena. +pub fn build(arena: &CallTraceArena) -> Vec { + let mut fst = EvmFoldedStackTraceBuilder::default(); + fst.process_call_node(arena.nodes(), 0); + fst.build() +} + +/// Wrapper for building a folded stack trace using EVM call trace node. +#[derive(Default)] +pub struct EvmFoldedStackTraceBuilder { + /// Raw folded stack trace builder. + fst: FoldedStackTraceBuilder, +} + +impl EvmFoldedStackTraceBuilder { + /// Returns the folded stack trace. + pub fn build(self) -> Vec { + self.fst.build() + } + + /// Creates an entry for a EVM CALL in the folded stack trace. This method recursively processes + /// all the children nodes of the call node and at the end it exits. + pub fn process_call_node(&mut self, nodes: &[CallTraceNode], idx: usize) { + let node = &nodes[idx]; + + let func_name = if node.trace.kind.is_any_create() { + let default_contract_name = "Contract".to_string(); + let contract_name = node.trace.decoded.label.as_ref().unwrap_or(&default_contract_name); + format!("new {contract_name}") + } else { + let selector = node + .selector() + .map(|selector| selector.encode_hex_with_prefix()) + .unwrap_or("fallback".to_string()); + let signature = + node.trace.decoded.call_data.as_ref().map(|dc| &dc.signature).unwrap_or(&selector); + + if let Some(label) = &node.trace.decoded.label { + format!("{label}.{signature}") + } else { + signature.clone() + } + }; + + self.fst.enter(func_name, node.trace.gas_used as i64); + + // Track internal function step exits to do in this call context. + let mut step_exits = vec![]; + + // Process children nodes. + for order in &node.ordering { + match order { + TraceMemberOrder::Call(child_idx) => { + let child_node_idx = node.children[*child_idx]; + self.process_call_node(nodes, child_node_idx); + } + TraceMemberOrder::Step(step_idx) => { + self.exit_previous_steps(&mut step_exits, *step_idx); + self.process_step(&node.trace.steps, *step_idx, &mut step_exits) + } + TraceMemberOrder::Log(_) => {} + } + } + + // Exit pending internal function calls if any. + for _ in 0..step_exits.len() { + self.fst.exit(); + } + + // Exit from this call context in the folded stack trace. + self.fst.exit(); + } + + /// Creates an entry for an internal function call in the folded stack trace. This method only + /// enters the function in the folded stack trace, we cannot exit since we need to exit at a + /// future step. Hence, we keep track of the step end index in the `step_exits`. + fn process_step( + &mut self, + steps: &[CallTraceStep], + step_idx: usize, + step_exits: &mut Vec, + ) { + let step = &steps[step_idx]; + if let Some(decoded_step) = &step.decoded { + match decoded_step { + DecodedTraceStep::InternalCall(decoded_internal_call, step_end_idx) => { + let gas_used = steps[*step_end_idx].gas_used.saturating_sub(step.gas_used); + self.fst.enter(decoded_internal_call.func_name.clone(), gas_used as i64); + step_exits.push(*step_end_idx); + } + DecodedTraceStep::Line(_) => {} + } + } + } + + /// Exits all the previous internal calls that should end before starting step_idx. + fn exit_previous_steps(&mut self, step_exits: &mut Vec, step_idx: usize) { + let initial_length = step_exits.len(); + step_exits.retain(|&number| number > step_idx); + + let num_exits = initial_length - step_exits.len(); + for _ in 0..num_exits { + self.fst.exit(); + } + } +} + +/// Helps to translate a function enter-exit flow into a folded stack trace. +/// +/// Example: +/// fn top() { child_a(); child_b() } // consumes 500 gas +/// fn child_a() {} // consumes 100 gas +/// fn child_b() {} // consumes 200 gas +/// +/// For execution of the `top` function looks like: +/// 1. enter `top` +/// 2. enter `child_a` +/// 3. exit `child_a` +/// 4. enter `child_b` +/// 5. exit `child_b` +/// 6. exit `top` +/// +/// The translated folded stack trace lines look like: +/// 1. top +/// 2. top;child_a +/// 3. top;child_b +/// +/// Including the gas consumed by the function by itself. +/// 1. top 200 // 500 - 100 - 200 +/// 2. top;child_a 100 +/// 3. top;child_b 200 +#[derive(Debug, Default)] +pub struct FoldedStackTraceBuilder { + /// Trace entries. + traces: Vec, + /// Number of exits to be done before entering a new function. + exits: usize, +} + +#[derive(Debug, Default)] +struct TraceEntry { + /// Names of all functions in the call stack of this trace. + names: Vec, + /// Gas consumed by this function, allowed to be negative due to refunds. + gas: i64, +} + +impl FoldedStackTraceBuilder { + /// Enter execution of a function call that consumes `gas`. + pub fn enter(&mut self, label: String, gas: i64) { + let mut names = self.traces.last().map(|entry| entry.names.clone()).unwrap_or_default(); + + while self.exits > 0 { + names.pop(); + self.exits -= 1; + } + + names.push(label); + self.traces.push(TraceEntry { names, gas }); + } + + /// Exit execution of a function call. + pub fn exit(&mut self) { + self.exits += 1; + } + + /// Returns folded stack trace. + pub fn build(mut self) -> Vec { + self.subtract_children(); + self.build_without_subtraction() + } + + /// Internal method to build the folded stack trace without subtracting gas consumed by + /// the children function calls. + fn build_without_subtraction(&mut self) -> Vec { + let mut lines = Vec::new(); + for TraceEntry { names, gas } in self.traces.iter() { + lines.push(format!("{} {}", names.join(";"), gas)); + } + lines + } + + /// Subtracts gas consumed by the children function calls from the parent function calls. + 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 TraceEntry { names, gas } = &right[0]; + if names.len() > 1 { + let parent_trace_to_match = &names[..names.len() - 1]; + for parent in left.iter_mut().rev() { + if parent.names == parent_trace_to_match { + parent.gas -= gas; + break; + } + } + } + } + } +} + +mod tests { + #[test] + fn test_fst_1() { + let mut trace = super::FoldedStackTraceBuilder::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.build_without_subtraction(), + vec![ + "top 500", // + "top;child_a 100", + "top;child_b 200", + ] + ); + assert_eq!( + trace.build(), + vec![ + "top 200", // 500 - 100 - 200 + "top;child_a 100", + "top;child_b 200", + ] + ); + } + + #[test] + fn test_fst_2() { + let mut trace = super::FoldedStackTraceBuilder::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.build_without_subtraction(), + vec![ + "top 500", // + "top;child_a 300", + "top;child_a;child_b 100", + "top;child_c 100", + ] + ); + + assert_eq!( + trace.build(), + 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_fst_3() { + let mut trace = super::FoldedStackTraceBuilder::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.build_without_subtraction(), + vec![ + "top 1700", // + "top;child_a 500", + "top;child_b 500", + "top;child_b;child_c 500", + "top2 1700", + ] + ); + + assert_eq!( + trace.build(), + 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 4a80b164ad33b..fb1665a5a6cd1 100644 --- a/crates/evm/traces/src/lib.rs +++ b/crates/evm/traces/src/lib.rs @@ -42,6 +42,8 @@ pub use decoder::{CallTraceDecoder, CallTraceDecoderBuilder}; pub mod debug; pub use debug::DebugTraceIdentifier; +pub mod folded_stack_trace; + pub type Traces = Vec<(TraceKind, SparsedTraceArena)>; /// Trace arena keeping track of ignored trace items. diff --git a/crates/forge/Cargo.toml b/crates/forge/Cargo.toml index 4f9110ff60e74..77034d6194de2 100644 --- a/crates/forge/Cargo.toml +++ b/crates/forge/Cargo.toml @@ -78,6 +78,7 @@ dialoguer = { version = "0.11", default-features = false } dunce.workspace = true futures.workspace = true indicatif = "0.17" +inferno = { version = "0.11.19", default-features = false } itertools.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 b0792a8cfc562..1cc647b0cb8af 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -1,7 +1,7 @@ use super::{install, test::filter::ProjectPathsAwareFilter, watch::WatchArgs}; use alloy_primitives::U256; use clap::{Parser, ValueHint}; -use eyre::Result; +use eyre::{Context, OptionExt, Result}; use forge::{ decode::decode_console_logs, gas_report::GasReport, @@ -9,7 +9,7 @@ use forge::{ result::{SuiteResult, TestOutcome, TestStatus}, traces::{ debug::{ContractSources, DebugTraceIdentifier}, - decode_trace_arena, + decode_trace_arena, folded_stack_trace, identifier::SignaturesIdentifier, render_trace_arena, CallTraceDecoderBuilder, InternalTraceMode, TraceKind, }, @@ -81,6 +81,14 @@ pub struct TestArgs { #[arg(long, value_name = "TEST_FUNCTION")] debug: Option, + /// Generate a flamegraph for a single test. Implies `--decode-internal`. + #[arg(long, conflicts_with = "flamechart")] + flamegraph: bool, + + /// Generate a flamechart for a single test. Implies `--decode-internal`. + #[arg(long, conflicts_with = "flamegraph")] + flamechart: bool, + /// Whether to identify internal functions in traces. /// /// If no argument is passed to this flag, it will trace internal functions scope and decode @@ -254,7 +262,7 @@ impl TestArgs { /// configured filter will be executed /// /// Returns the test results for all matching tests. - pub async fn execute_tests(self) -> Result { + pub async fn execute_tests(mut self) -> Result { // Merge all configs. let (mut config, mut evm_opts) = self.load_config_and_evm_opts_emit_warnings()?; @@ -308,14 +316,22 @@ impl TestArgs { .profiles(profiles) .build(&output, project_root)?; + let should_debug = self.debug.is_some(); + let should_draw = self.flamegraph || self.flamechart; + // 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) || self.flamegraph || self.flamechart { evm_opts.verbosity = 3; } let env = evm_opts.evm_env().await?; + // Enable internal tracing for more informative flamegraph. + if should_draw { + self.decode_internal = Some(None); + } + // Choose the internal function tracing mode, if --decode-internal is provided. let decode_internal = if let Some(maybe_fn) = self.decode_internal.as_ref() { if maybe_fn.is_some() { @@ -330,7 +346,6 @@ impl TestArgs { }; // 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) @@ -366,18 +381,60 @@ impl TestArgs { )?; let libraries = runner.libraries.clone(); - let outcome = self.run_tests(runner, config, verbosity, &filter, &output).await?; + let mut outcome = self.run_tests(runner, config, verbosity, &filter, &output).await?; + + if should_draw { + let (suite_name, test_name, mut test_result) = + outcome.remove_first().ok_or_eyre("no tests were executed")?; + + let arena = test_result + .traces + .iter_mut() + .find_map( + |(kind, arena)| { + if *kind == TraceKind::Execution { + Some(arena) + } else { + None + } + }, + ) + .unwrap(); + + // Decode traces. + let decoder = outcome.last_run_decoder.as_ref().unwrap(); + decode_trace_arena(arena, decoder).await?; + let mut fst = folded_stack_trace::build(arena); + + let label = if self.flamegraph { "flamegraph" } else { "flamechart" }; + let contract = suite_name.split(':').last().unwrap(); + let test_name = test_name.trim_end_matches("()"); + let file_name = format!("cache/{label}_{contract}_{test_name}.svg"); + let file = std::fs::File::create(&file_name).wrap_err("failed to create file")?; + + let mut options = inferno::flamegraph::Options::default(); + options.title = format!("{label} {contract}::{test_name}"); + options.count_name = "gas".to_string(); + if self.flamechart { + options.flame_chart = true; + fst.reverse(); + } + + // Generate SVG. + inferno::flamegraph::from_lines(&mut options, fst.iter().map(|s| s.as_str()), file) + .wrap_err("failed to write svg")?; + println!("\nSaved to {file_name}"); + + // Open SVG in default program. + if opener::open(&file_name).is_err() { + println!("\nFailed to open {file_name}. Please open it manually."); + } + } if should_debug { // Get first non-empty suite result. We will have only one such entry. - let Some((_, test_result)) = outcome - .results - .iter() - .find(|(_, r)| !r.test_results.is_empty()) - .map(|(_, r)| (r, r.test_results.values().next().unwrap())) - else { - return Err(eyre::eyre!("no tests were executed")); - }; + let (_, _, test_result) = + outcome.remove_first().ok_or_eyre("no tests were executed")?; let sources = ContractSources::from_project_output(&output, project.root(), Some(&libraries))?; @@ -417,13 +474,17 @@ 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.as_ref().map_or(false, |v| v.is_some())) && + if (self.debug.is_some() || + self.decode_internal.as_ref().map_or(false, |v| v.is_some()) || + self.flamegraph || + self.flamechart) && 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.flamegraph {"generate a flamegraph"} else if self.flamechart {"generate a flamechart"} else {"run the debugger"}, ); } @@ -489,7 +550,11 @@ 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 || + self.flamechart; // Print suite header. println!(); diff --git a/crates/forge/src/result.rs b/crates/forge/src/result.rs index e7953575f3b50..aa88bc4016fcf 100644 --- a/crates/forge/src/result.rs +++ b/crates/forge/src/result.rs @@ -184,6 +184,18 @@ impl TestOutcome { // TODO: Avoid process::exit std::process::exit(1); } + + /// Removes first test result, if any. + pub fn remove_first(&mut self) -> Option<(String, String, TestResult)> { + self.results.iter_mut().find_map(|(suite_name, suite)| { + if let Some(test_name) = suite.test_results.keys().next().cloned() { + let result = suite.test_results.remove(&test_name).unwrap(); + Some((suite_name.clone(), test_name, result)) + } else { + None + } + }) + } } /// A set of test results for a single test suite, which is all the tests in a single contract. diff --git a/deny.toml b/deny.toml index 18fcd55652da0..be64e55e32b42 100644 --- a/deny.toml +++ b/deny.toml @@ -50,6 +50,7 @@ allow = [ "BSL-1.0", "0BSD", "MPL-2.0", + "CDDL-1.0", ] # Allow 1 or more licenses on a per-crate basis, so that particular licenses