feat(test): include data from fuzz/invariant runs in gas reports

crates/config/src/fuzz.rs

@@ -22,6 +22,8 @@ pub struct FuzzConfig {
     /// The fuzz dictionary configuration
     #[serde(flatten)]
     pub dictionary: FuzzDictionaryConfig,
+    /// Number of runs to execute and include in the gas report.
+    pub gas_report_samples: u32,
 }
 
 impl Default for FuzzConfig {
@@ -31,6 +33,7 @@ impl Default for FuzzConfig {
             max_test_rejects: 65536,
             seed: None,
             dictionary: FuzzDictionaryConfig::default(),
+            gas_report_samples: 256,
         }
     }
 }

crates/config/src/invariant.rs

@@ -35,6 +35,8 @@ pub struct InvariantConfig {
     /// The maximum number of rejects via `vm.assume` which can be encountered during a single
     /// invariant run.
     pub max_assume_rejects: u32,
+    /// Number of runs to execute and include in the gas report.
+    pub gas_report_samples: u32,
 }
 
 impl Default for InvariantConfig {
@@ -49,6 +51,7 @@ impl Default for InvariantConfig {
             shrink_run_limit: 2usize.pow(18_u32),
             preserve_state: false,
             max_assume_rejects: 65536,
+            gas_report_samples: 256,
         }
     }
 }

crates/evm/evm/src/executors/fuzz/mod.rs

@@ -71,8 +71,11 @@ impl FuzzedExecutor {
         // Stores the result and calldata of the last failed call, if any.
         let counterexample: RefCell<(Bytes, RawCallResult)> = RefCell::default();
 
-        // Stores the last successful call trace
-        let traces: RefCell<Option<CallTraceArena>> = RefCell::default();
+        // We want to collect at least one trace which will be displayed to user.
+        let max_traces_to_collect = std::cmp::max(1, self.config.gas_report_samples) as usize;
+
+        // Stores up to `max_traces_to_collect` traces.
+        let traces: RefCell<Vec<CallTraceArena>> = RefCell::default();
 
         // Stores coverage information for all fuzz cases
         let coverage: RefCell<Option<HitMaps>> = RefCell::default();
@@ -103,8 +106,12 @@ impl FuzzedExecutor {
                     if first_case.is_none() {
                         first_case.replace(case.case);
                     }
-
-                    traces.replace(case.traces);
+                    if let Some(call_traces) = case.traces {
+                        if traces.borrow().len() == max_traces_to_collect {
+                            traces.borrow_mut().pop();
+                        }
+                        traces.borrow_mut().push(call_traces);
+                    }
 
                     if let Some(prev) = coverage.take() {
                         // Safety: If `Option::or` evaluates to `Some`, then `call.coverage` must
@@ -137,6 +144,10 @@ impl FuzzedExecutor {
         });
 
         let (calldata, call) = counterexample.into_inner();
+
+        let mut traces = traces.into_inner();
+        let last_run_traces = if run_result.is_ok() { traces.pop() } else { call.traces.clone() };
+
         let mut result = FuzzTestResult {
             first_case: first_case.take().unwrap_or_default(),
             gas_by_case: gas_by_case.take(),
@@ -146,7 +157,8 @@ impl FuzzedExecutor {
             decoded_logs: decode_console_logs(&call.logs),
             logs: call.logs,
             labeled_addresses: call.labels,
-            traces: if run_result.is_ok() { traces.into_inner() } else { call.traces.clone() },
+            traces: last_run_traces,
+            gas_report_traces: traces,
             coverage: coverage.into_inner(),
         };
 

crates/evm/evm/src/executors/invariant/error.rs

@@ -50,6 +50,9 @@ pub struct InvariantFuzzTestResult {
     /// The entire inputs of the last run of the invariant campaign, used for
     /// replaying the run for collecting traces.
     pub last_run_inputs: Vec<BasicTxDetails>,
+
+    /// Additional traces used for gas report construction.
+    pub gas_report_traces: Vec<Vec<CallTraceArena>>,
 }
 
 #[derive(Clone, Debug)]

crates/evm/evm/src/executors/invariant/mod.rs

@@ -23,6 +23,7 @@ use foundry_evm_fuzz::{
     },
     FuzzCase, FuzzedCases,
 };
+use foundry_evm_traces::CallTraceArena;
 use parking_lot::{Mutex, RwLock};
 use proptest::{
     strategy::{BoxedStrategy, Strategy, ValueTree},
@@ -123,6 +124,9 @@ impl<'a> InvariantExecutor<'a> {
         // Stores the calldata in the last run.
         let last_run_calldata: RefCell<Vec<BasicTxDetails>> = RefCell::new(vec![]);
 
+        // Stores additional traces for gas report.
+        let gas_report_traces: RefCell<Vec<Vec<CallTraceArena>>> = RefCell::default();
+
         // Let's make sure the invariant is sound before actually starting the run:
         // We'll assert the invariant in its initial state, and if it fails, we'll
         // already know if we can early exit the invariant run.
@@ -162,6 +166,9 @@ impl<'a> InvariantExecutor<'a> {
             // Created contracts during a run.
             let mut created_contracts = vec![];
 
+            // Traces of each call of the sequence.
+            let mut run_traces = Vec::new();
+
             let mut current_run = 0;
             let mut assume_rejects_counter = 0;
 
@@ -233,6 +240,7 @@ impl<'a> InvariantExecutor<'a> {
                             self.config.fail_on_revert,
                             self.config.shrink_sequence,
                             self.config.shrink_run_limit,
+                            &mut run_traces,
                         );
 
                     if !can_continue || current_run == self.config.depth - 1 {
@@ -265,6 +273,9 @@ impl<'a> InvariantExecutor<'a> {
                 }
             }
 
+            if gas_report_traces.borrow().len() < self.config.gas_report_samples as usize {
+                gas_report_traces.borrow_mut().push(run_traces);
+            }
             fuzz_cases.borrow_mut().push(FuzzedCases::new(fuzz_runs));
 
             Ok(())
@@ -280,6 +291,7 @@ impl<'a> InvariantExecutor<'a> {
             cases: fuzz_cases.into_inner(),
             reverts,
             last_run_inputs: last_run_calldata.take(),
+            gas_report_traces: gas_report_traces.into_inner(),
         })
     }
 
@@ -764,6 +776,7 @@ fn can_continue(
     fail_on_revert: bool,
     shrink_sequence: bool,
     shrink_run_limit: usize,
+    run_traces: &mut Vec<CallTraceArena>,
 ) -> RichInvariantResults {
     let mut call_results = None;
 
@@ -775,6 +788,10 @@ fn can_continue(
 
     // Assert invariants IFF the call did not revert and the handlers did not fail.
     if !call_result.reverted && !handlers_failed {
+        if let Some(traces) = call_result.traces {
+            run_traces.push(traces);
+        }
+
         call_results = assert_invariants(
             invariant_contract,
             executor,

crates/evm/fuzz/src/lib.rs

@@ -150,6 +150,10 @@ pub struct FuzzTestResult {
     /// `num(fuzz_cases)` traces, one for each run, which is neither helpful nor performant.
     pub traces: Option<CallTraceArena>,
 
+    /// Additional traces used for gas report construction.
+    /// Those traces should not be displayed.
+    pub gas_report_traces: Vec<CallTraceArena>,
+
     /// Raw coverage info
     pub coverage: Option<HitMaps>,
 }

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

@@ -145,6 +145,10 @@ impl TestArgs {
         // Explicitly enable isolation for gas reports for more correct gas accounting
         if self.gas_report {
             evm_opts.isolate = true;
+        } else {
+            // Do not collect gas report traces if gas report is not enabled.
+            config.fuzz.gas_report_samples = 0;
+            config.invariant.gas_report_samples = 0;
         }
 
         // Set up the project.
@@ -411,7 +415,26 @@ impl TestArgs {
                 }
 
                 if let Some(gas_report) = &mut gas_report {
-                    gas_report.analyze(&result.traces, &decoder).await;
+                    gas_report
+                        .analyze(result.traces.iter().map(|(_, arena)| arena), &decoder)
+                        .await;
+
+                    for trace in result.gas_report_traces.iter() {
+                        decoder.clear_addresses();
+
+                        // Re-execute setup and deployment traces to collect identities created in
+                        // setUp and constructor.
+                        for (kind, arena) in &result.traces {
+                            if !matches!(kind, TraceKind::Execution) {
+                                decoder.identify(arena, &mut local_identifier);
+                            }
+                        }
+
+                        for arena in trace {
+                            decoder.identify(arena, &mut local_identifier);
+                            gas_report.analyze([arena], &decoder).await;
+                        }
+                    }
                 }
             }
 
@@ -433,7 +456,9 @@ impl TestArgs {
         outcome.decoder = Some(decoder);
 
         if let Some(gas_report) = gas_report {
-            shell::println(gas_report.finalize())?;
+            let finalized = gas_report.finalize();
+            shell::println(&finalized)?;
+            outcome.gas_report = Some(finalized);
         }
 
         if !outcome.results.is_empty() {
@@ -530,6 +555,7 @@ fn list(
 mod tests {
     use super::*;
     use foundry_config::Chain;
+    use foundry_test_utils::forgetest_async;
 
     #[test]
     fn watch_parse() {
@@ -563,4 +589,62 @@ mod tests {
         test("--chain-id=1", Chain::mainnet());
         test("--chain-id=42", Chain::from_id(42));
     }
+
+    forgetest_async!(gas_report_fuzz_invariant, |prj, _cmd| {
+        prj.insert_ds_test();
+        prj.add_source(
+            "Contracts.sol",
+            r#"
+//SPDX-license-identifier: MIT
+
+import "./test.sol";
+
+contract Foo {
+    function foo() public {}
+}
+
+contract Bar {
+    function bar() public {}
+}
+
+
+contract FooBarTest is DSTest {
+    Foo public targetContract;
+
+    function setUp() public {
+        targetContract = new Foo();
+    }
+
+    function invariant_dummy() public {
+        assertTrue(true);
+    }
+
+    function testFuzz_bar(uint256 _val) public {
+        (new Bar()).bar();
+    }
+}
+        "#,
+        )
+        .unwrap();
+
+        let args = TestArgs::parse_from([
+            "foundry-cli",
+            "--gas-report",
+            "--root",
+            &prj.root().to_string_lossy(),
+            "--silent",
+        ]);
+
+        let outcome = args.run().await.unwrap();
+        let gas_report = outcome.gas_report.unwrap();
+
+        assert_eq!(gas_report.contracts.len(), 3);
+        let call_cnts = gas_report
+            .contracts
+            .values()
+            .flat_map(|c| c.functions.values().flat_map(|f| f.values().map(|v| v.calls.len())))
+            .collect::<Vec<_>>();
+        // assert that all functions were called at least 100 times
+        assert!(call_cnts.iter().all(|c| *c > 100));
+    });
 }

crates/forge/src/gas_report.rs

@@ -3,7 +3,7 @@
 use crate::{
     constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS},
     hashbrown::HashSet,
-    traces::{CallTraceArena, CallTraceDecoder, CallTraceNode, DecodedCallData, TraceKind},
+    traces::{CallTraceArena, CallTraceDecoder, CallTraceNode, DecodedCallData},
 };
 use comfy_table::{presets::ASCII_MARKDOWN, *};
 use foundry_common::{calc, TestFunctionExt};
@@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize};
 use std::{collections::BTreeMap, fmt::Display};
 
 /// Represents the gas report for a set of contracts.
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct GasReport {
     /// Whether to report any contracts.
     report_any: bool,
@@ -22,7 +22,7 @@ pub struct GasReport {
     ignore: HashSet<String>,
     /// All contracts that were analyzed grouped by their identifier
     /// ``test/Counter.t.sol:CounterTest
-    contracts: BTreeMap<String, ContractInfo>,
+    pub contracts: BTreeMap<String, ContractInfo>,
 }
 
 impl GasReport {
@@ -61,10 +61,10 @@ impl GasReport {
     /// Analyzes the given traces and generates a gas report.
     pub async fn analyze(
         &mut self,
-        traces: &[(TraceKind, CallTraceArena)],
+        arenas: impl IntoIterator<Item = &CallTraceArena>,
         decoder: &CallTraceDecoder,
     ) {
-        for node in traces.iter().flat_map(|(_, arena)| arena.nodes()) {
+        for node in arenas.into_iter().flat_map(|arena| arena.nodes()) {
             self.analyze_node(node, decoder).await;
         }
     }
@@ -78,7 +78,7 @@ impl GasReport {
 
         // Only include top-level calls which accout for calldata and base (21.000) cost.
         // Only include Calls and Creates as only these calls are isolated in inspector.
-        if trace.depth != 1 &&
+        if trace.depth > 1 &&
             (trace.kind == CallKind::Call ||
                 trace.kind == CallKind::Create ||
                 trace.kind == CallKind::Create2)
@@ -186,15 +186,15 @@ impl Display for GasReport {
     }
 }
 
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct ContractInfo {
     pub gas: u64,
     pub size: usize,
     /// Function name -> Function signature -> GasInfo
     pub functions: BTreeMap<String, BTreeMap<String, GasInfo>>,
 }
 
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct GasInfo {
     pub calls: Vec<u64>,
     pub min: u64,