feat(forge): add gas reports for tests

cli/src/cmd/test.rs

@@ -7,7 +7,9 @@ use crate::{
 use ansi_term::Colour;
 use clap::{AppSettings, Parser};
 use ethers::solc::{ArtifactOutput, Project};
-use evm_adapters::{call_tracing::ExecutionInfo, evm_opts::EvmOpts, sputnik::helpers::vm};
+use evm_adapters::{
+    call_tracing::ExecutionInfo, evm_opts::EvmOpts, gas_report::GasReport, sputnik::helpers::vm,
+};
 use forge::{MultiContractRunnerBuilder, TestFilter};
 use foundry_config::{figment::Figment, Config};
 use std::collections::BTreeMap;
@@ -23,27 +25,31 @@ pub struct Filter {
 
     #[clap(
         long = "match-test",
+        alias = "mt",
         help = "only run test methods matching regex",
         conflicts_with = "pattern"
     )]
     test_pattern: Option<regex::Regex>,
 
     #[clap(
         long = "no-match-test",
+        alias = "nmt",
         help = "only run test methods not matching regex",
         conflicts_with = "pattern"
     )]
     test_pattern_inverse: Option<regex::Regex>,
 
     #[clap(
         long = "match-contract",
+        alias = "mc",
         help = "only run test methods in contracts matching regex",
         conflicts_with = "pattern"
     )]
     contract_pattern: Option<regex::Regex>,
 
     #[clap(
         long = "no-match-contract",
+        alias = "nmc",
         help = "only run test methods in contracts not matching regex",
         conflicts_with = "pattern"
     )]
@@ -88,6 +94,9 @@ pub struct TestArgs {
     #[clap(help = "print the test results in json format", long, short)]
     json: bool,
 
+    #[clap(help = "print a gas report", long = "gas-report")]
+    gas_report: bool,
+
     #[clap(flatten)]
     evm_opts: EvmArgs,
 
@@ -138,7 +147,15 @@ impl Cmd for TestArgs {
             .evm_cfg(evm_cfg)
             .sender(evm_opts.sender);
 
-        test(builder, project, evm_opts, filter, json, allow_failure)
+        test(
+            builder,
+            project,
+            evm_opts,
+            filter,
+            json,
+            allow_failure,
+            (self.gas_report, config.gas_reports),
+        )
     }
 }
 
@@ -251,16 +268,26 @@ fn short_test_result(name: &str, result: &forge::TestResult) {
 fn test<A: ArtifactOutput + 'static>(
     builder: MultiContractRunnerBuilder,
     project: Project<A>,
-    evm_opts: EvmOpts,
+    mut evm_opts: EvmOpts,
     filter: Filter,
     json: bool,
     allow_failure: bool,
+    gas_reports: (bool, Vec<String>),
 ) -> eyre::Result<TestOutcome> {
     let verbosity = evm_opts.verbosity;
+    let gas_reporting = gas_reports.0;
+
+    if gas_reporting && evm_opts.verbosity < 3 {
+        // force evm to do tracing, but dont hit the verbosity print path
+        evm_opts.verbosity = 3;
+    }
+
     let mut runner = builder.build(project, evm_opts)?;
 
     let results = runner.test(&filter)?;
 
+    let mut gas_report = GasReport::new(gas_reports.1);
+
     let (funcs, events, errors) = runner.execution_info;
     if json {
         let res = serde_json::to_string(&results)?;
@@ -277,6 +304,15 @@ fn test<A: ArtifactOutput + 'static>(
             }
 
             for (name, result) in tests {
+                // build up gas report
+                if gas_reporting {
+                    if let (Some(traces), Some(identified_contracts)) =
+                        (&result.traces, &result.identified_contracts)
+                    {
+                        gas_report.analyze(traces, identified_contracts);
+                    }
+                }
+
                 short_test_result(name, result);
 
                 // adds a linebreak only if there were any traces or logs, so that the
@@ -339,5 +375,10 @@ fn test<A: ArtifactOutput + 'static>(
         }
     }
 
+    if gas_reporting {
+        gas_report.finalize();
+        println!("{}", gas_report);
+    }
+
     Ok(TestOutcome::new(results, allow_failure))
 }

config/README.md

@@ -69,6 +69,7 @@ libraries = []
 cache = true
 force = false
 evm_version = 'london'
+gas_reports = ["*"]
 ## Sets the concrete solc version to use, this overrides the `auto_detect_solc` value
 # solc_version = '0.8.10'
 auto_detect_solc = true

config/src/lib.rs

@@ -91,6 +91,8 @@ pub struct Config {
     /// evm version to use
     #[serde(with = "from_str_lowercase")]
     pub evm_version: EvmVersion,
+    /// list of contracts to report gas of
+    pub gas_reports: Vec<String>,
     /// Concrete solc version to use if any.
     ///
     /// This takes precedence over `auto_detect_solc`, if a version is set then this overrides
@@ -638,6 +640,7 @@ impl Default for Config {
             cache: true,
             force: false,
             evm_version: Default::default(),
+            gas_reports: vec!["*".to_string()],
             solc_version: None,
             auto_detect_solc: true,
             optimizer: true,

evm-adapters/Cargo.toml

@@ -29,6 +29,7 @@ revm_precompiles = { git = "https://github.com/bluealloy/revm", default-features
 serde_json = "1.0.72"
 serde = "1.0.130"
 ansi_term = "0.12.1"
+comfy-table = "5.0.0"
 
 [dev-dependencies]
 evmodin = { git = "https://github.com/vorot93/evmodin", features = ["util"] }

evm-adapters/src/gas_report.rs

@@ -0,0 +1,167 @@
+use crate::CallTraceArena;
+use ethers::{
+    abi::Abi,
+    types::{H160, U256},
+};
+use serde::{Deserialize, Serialize};
+use std::{collections::BTreeMap, fmt::Display};
+
+#[cfg(feature = "sputnik")]
+use crate::sputnik::cheatcodes::cheatcode_handler::{CHEATCODE_ADDRESS, CONSOLE_ADDRESS};
+
+use comfy_table::{modifiers::UTF8_ROUND_CORNERS, presets::UTF8_FULL, *};
+
+#[derive(Default, Debug, Serialize, Deserialize)]
+pub struct GasReport {
+    pub report_for: Vec<String>,
+    pub contracts: BTreeMap<String, ContractInfo>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Default)]
+pub struct ContractInfo {
+    pub gas: U256,
+    pub size: U256,
+    pub functions: BTreeMap<String, GasInfo>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Default)]
+pub struct GasInfo {
+    pub calls: Vec<U256>,
+    pub min: U256,
+    pub mean: U256,
+    pub median: U256,
+    pub max: U256,
+}
+
+impl GasReport {
+    pub fn new(report_for: Vec<String>) -> Self {
+        Self { report_for, ..Default::default() }
+    }
+
+    pub fn analyze(
+        &mut self,
+        traces: &[CallTraceArena],
+        identified_contracts: &BTreeMap<H160, (String, Abi)>,
+    ) {
+        let report_for_all = self.report_for.is_empty() || self.report_for.iter().any(|s| s == "*");
+        traces.iter().for_each(|trace| {
+            self.analyze_trace(trace, identified_contracts, report_for_all);
+        });
+    }
+
+    fn analyze_trace(
+        &mut self,
+        trace: &CallTraceArena,
+        identified_contracts: &BTreeMap<H160, (String, Abi)>,
+        report_for_all: bool,
+    ) {
+        self.analyze_node(trace.entry, trace, identified_contracts, report_for_all);
+    }
+
+    fn analyze_node(
+        &mut self,
+        node_index: usize,
+        arena: &CallTraceArena,
+        identified_contracts: &BTreeMap<H160, (String, Abi)>,
+        report_for_all: bool,
+    ) {
+        let node = &arena.arena[node_index];
+        let trace = &node.trace;
+
+        #[cfg(feature = "sputnik")]
+        if trace.addr == *CHEATCODE_ADDRESS || trace.addr == *CONSOLE_ADDRESS {
+            return
+        }
+
+        if let Some((name, abi)) = identified_contracts.get(&trace.addr) {
+            let report_for = self.report_for.iter().any(|s| s == name);
+            if !report_for && abi.functions().any(|func| func.name == "IS_TEST") {
+                // do nothing
+            } else if report_for || report_for_all {
+                // report for this contract
+                let mut contract =
+                    self.contracts.entry(name.to_string()).or_insert_with(Default::default);
+
+                if trace.created {
+                    contract.gas = trace.cost.into();
+                    contract.size = trace.data.len().into();
+                } else if trace.data.len() >= 4 {
+                    let func =
+                        abi.functions().find(|func| func.short_signature() == trace.data[0..4]);
+
+                    if let Some(func) = func {
+                        let function = contract
+                            .functions
+                            .entry(func.name.clone())
+                            .or_insert_with(Default::default);
+                        function.calls.push(trace.cost.into());
+                    }
+                }
+            }
+        }
+        node.children.iter().for_each(|index| {
+            self.analyze_node(*index, arena, identified_contracts, report_for_all);
+        });
+    }
+
+    pub fn finalize(&mut self) {
+        self.contracts.iter_mut().for_each(|(_name, contract)| {
+            contract.functions.iter_mut().for_each(|(_name, func)| {
+                func.calls.sort();
+                func.min = func.calls.first().cloned().unwrap_or_default();
+                func.max = func.calls.last().cloned().unwrap_or_default();
+                func.mean =
+                    func.calls.iter().fold(U256::zero(), |acc, x| acc + x) / func.calls.len();
+
+                let len = func.calls.len();
+                func.median = if len > 0 {
+                    if len % 2 == 0 {
+                        (func.calls[len / 2 - 1] + func.calls[len / 2]) / 2
+                    } else {
+                        func.calls[len / 2]
+                    }
+                } else {
+                    0.into()
+                };
+            });
+        });
+    }
+}
+
+impl Display for GasReport {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
+        for (name, contract) in self.contracts.iter() {
+            let mut table = Table::new();
+            table.load_preset(UTF8_FULL).apply_modifier(UTF8_ROUND_CORNERS);
+            table.set_header(vec![Cell::new(format!("{} contract", name))
+                .add_attribute(Attribute::Bold)
+                .fg(Color::Green)]);
+            table.add_row(vec![
+                Cell::new("Deployment Cost").add_attribute(Attribute::Bold).fg(Color::Cyan),
+                Cell::new("Deployment Size").add_attribute(Attribute::Bold).fg(Color::Cyan),
+            ]);
+            table.add_row(vec![contract.gas.to_string(), contract.size.to_string()]);
+
+            table.add_row(vec![
+                Cell::new("Function Name").add_attribute(Attribute::Bold).fg(Color::Magenta),
+                Cell::new("min").add_attribute(Attribute::Bold).fg(Color::Green),
+                Cell::new("avg").add_attribute(Attribute::Bold).fg(Color::Yellow),
+                Cell::new("median").add_attribute(Attribute::Bold).fg(Color::Yellow),
+                Cell::new("max").add_attribute(Attribute::Bold).fg(Color::Red),
+                Cell::new("# calls").add_attribute(Attribute::Bold),
+            ]);
+            contract.functions.iter().for_each(|(fname, function)| {
+                table.add_row(vec![
+                    Cell::new(fname.to_string()).add_attribute(Attribute::Bold),
+                    Cell::new(function.min.to_string()).fg(Color::Green),
+                    Cell::new(function.mean.to_string()).fg(Color::Yellow),
+                    Cell::new(function.median.to_string()).fg(Color::Yellow),
+                    Cell::new(function.max.to_string()).fg(Color::Red),
+                    Cell::new(function.calls.len().to_string()),
+                ]);
+            });
+            writeln!(f, "{}", table)?
+        }
+        Ok(())
+    }
+}

evm-adapters/src/lib.rs

@@ -18,6 +18,8 @@ pub mod fuzz;
 
 pub mod call_tracing;
 
+pub mod gas_report;
+
 /// Helpers for easily constructing EVM objects.
 pub mod evm_opts;