Skip to main content

forge/
gas_report.rs

1//! Gas reports.
2
3use crate::{
4    constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS},
5    traces::{CallTraceArena, CallTraceDecoder, CallTraceNode, DecodedCallData},
6};
7use alloy_primitives::{Address, map::HashSet};
8use comfy_table::{
9    Cell, CellAlignment, Color, Table,
10    presets::{ASCII_FULL, ASCII_MARKDOWN},
11};
12use foundry_common::{TestFunctionExt, calc, get_contract_name, shell};
13use foundry_evm::traces::CallKind;
14use serde::{Deserialize, Serialize};
15use serde_json::json;
16use std::{collections::BTreeMap, fmt::Display};
17
18/// Represents the gas report for a set of contracts.
19#[derive(Clone, Debug, Default, Serialize, Deserialize)]
20pub struct GasReport {
21    /// Whether to report any contracts.
22    report_any: bool,
23    /// Contracts to generate the report for.
24    report_for: HashSet<String>,
25    /// Contracts to ignore when generating the report.
26    ignore: HashSet<String>,
27    /// Whether to include gas reports for tests.
28    include_tests: bool,
29    /// Additional network-specific cheatcode addresses omitted from reports.
30    #[serde(skip)]
31    extra_cheatcode_addresses: HashSet<Address>,
32    /// All contracts that were analyzed grouped by their identifier
33    /// ``test/Counter.t.sol:CounterTest
34    pub contracts: BTreeMap<String, ContractInfo>,
35}
36
37impl GasReport {
38    pub fn new(
39        report_for: impl IntoIterator<Item = String>,
40        ignore: impl IntoIterator<Item = String>,
41        include_tests: bool,
42        extra_cheatcode_addresses: impl IntoIterator<Item = Address>,
43    ) -> Self {
44        let report_for = report_for.into_iter().collect::<HashSet<_>>();
45        let report_any = report_for.is_empty() || report_for.contains("*");
46        Self {
47            report_any,
48            report_for,
49            ignore: ignore.into_iter().collect(),
50            include_tests,
51            extra_cheatcode_addresses: extra_cheatcode_addresses.into_iter().collect(),
52            contracts: BTreeMap::new(),
53        }
54    }
55
56    /// Whether the given contract should be reported.
57    #[instrument(level = "trace", skip(self), ret)]
58    fn should_report(&self, contract_name: &str) -> bool {
59        if self.ignore.contains(contract_name) {
60            let contains_anyway = self.report_for.contains(contract_name);
61            if contains_anyway {
62                // If the user listed the contract in 'gas_reports' (the foundry.toml field) a
63                // report for the contract is generated even if it's listed in the ignore
64                // list. This is addressed this way because getting a report you don't expect is
65                // preferable than not getting one you expect. A warning is printed to stderr
66                // indicating the "double listing".
67                let _ = sh_warn!(
68                    "{contract_name} is listed in both 'gas_reports' and 'gas_reports_ignore'."
69                );
70            }
71            return contains_anyway;
72        }
73        self.report_any || self.report_for.contains(contract_name)
74    }
75
76    fn is_internal_address(&self, address: Address) -> bool {
77        address == CHEATCODE_ADDRESS
78            || address == HARDHAT_CONSOLE_ADDRESS
79            || self.extra_cheatcode_addresses.contains(&address)
80    }
81
82    /// Analyzes the given traces and generates a gas report.
83    pub async fn analyze(
84        &mut self,
85        arenas: impl IntoIterator<Item = &CallTraceArena>,
86        decoder: &CallTraceDecoder,
87    ) {
88        for node in arenas.into_iter().flat_map(|arena| arena.nodes()) {
89            self.analyze_node(node, decoder).await;
90        }
91    }
92
93    async fn analyze_node(&mut self, node: &CallTraceNode, decoder: &CallTraceDecoder) {
94        let trace = &node.trace;
95        if self.is_internal_address(trace.address) {
96            return;
97        }
98        let Some(name) = decoder.contracts.get(&trace.address) else { return };
99        let contract_name = get_contract_name(name);
100        if !self.should_report(contract_name) {
101            return;
102        }
103
104        let contract_info = self.contracts.entry(name.clone()).or_default();
105        let is_create_call = trace.kind.is_any_create();
106        if is_create_call {
107            trace!(contract_name, "adding create size info");
108            contract_info.size = trace.data.len();
109        }
110
111        // Only include top-level calls which account for calldata and base (21.000) cost.
112        // Only include Calls and Creates as only these calls are isolated in inspector.
113        if trace.depth > 1 && (trace.kind == CallKind::Call || is_create_call) {
114            return;
115        }
116
117        if is_create_call {
118            trace!(contract_name, "adding create gas info");
119            contract_info.gas = trace.gas_used;
120        } else if let Some(DecodedCallData { signature, .. }) =
121            decoder.decode_function(trace).await.call_data
122        {
123            let name = signature.split('(').next().unwrap();
124            // Ignore any test/setup functions.
125            if self.include_tests || !name.test_function_kind().is_known() {
126                trace!(contract_name, signature, "adding gas info");
127                contract_info
128                    .functions
129                    .entry(name.to_string())
130                    .or_default()
131                    .entry(signature)
132                    .or_default()
133                    .frames
134                    .push(trace.gas_used);
135            }
136        }
137    }
138
139    /// Finalizes the gas report by calculating the min, max, mean, and median for each function.
140    #[must_use]
141    pub fn finalize(mut self) -> Self {
142        trace!("finalizing gas report");
143        for func in self
144            .contracts
145            .values_mut()
146            .flat_map(|c| c.functions.values_mut().flat_map(|s| s.values_mut()))
147        {
148            func.frames.sort_unstable();
149            func.min = func.frames.first().copied().unwrap_or_default();
150            func.max = func.frames.last().copied().unwrap_or_default();
151            func.mean = calc::mean(&func.frames);
152            func.median = calc::median_sorted(&func.frames);
153            func.calls = func.frames.len() as u64;
154        }
155        self
156    }
157
158    /// Contracts with at least one recorded function call.
159    fn reported_contracts(&self) -> impl Iterator<Item = (&String, &ContractInfo)> {
160        self.contracts.iter().filter(|(name, contract)| {
161            if contract.functions.is_empty() {
162                trace!(name, "gas report contract without functions");
163            }
164            !contract.functions.is_empty()
165        })
166    }
167}
168
169impl Display for GasReport {
170    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
171        if shell::is_json() {
172            return writeln!(f, "{}", self.format_json_output());
173        }
174        for (name, contract) in self.reported_contracts() {
175            writeln!(f, "\n{}", format_table_output(contract, name))?;
176        }
177        Ok(())
178    }
179}
180
181impl GasReport {
182    fn format_json_output(&self) -> String {
183        let contracts = self
184            .reported_contracts()
185            .map(|(name, contract)| {
186                let functions = contract
187                    .functions
188                    .values()
189                    .flat_map(|sigs| sigs.iter().map(|(sig, info)| (sig.replace(':', ""), info)))
190                    .collect::<BTreeMap<_, _>>();
191                json!({
192                    "contract": name,
193                    "deployment": { "gas": contract.gas, "size": contract.size },
194                    "functions": functions,
195                })
196            })
197            .collect::<Vec<_>>();
198        serde_json::to_string(&contracts).unwrap()
199    }
200}
201
202fn format_table_output(contract: &ContractInfo, name: &str) -> Table {
203    let num = |value: &dyn Display, color: Option<Color>| {
204        let cell = Cell::new(value.to_string()).set_alignment(CellAlignment::Right);
205        match color {
206            Some(color) => cell.fg(color),
207            None => cell,
208        }
209    };
210
211    let mut table = Table::new();
212    if shell::is_markdown() {
213        table.load_style(ASCII_MARKDOWN);
214    } else {
215        table.load_style(ASCII_FULL.with_rounded_corners());
216    }
217    table.set_header(vec![Cell::new(format!("{name} Contract")).fg(Color::Magenta)]);
218    table.add_row(vec![
219        Cell::new("Deployment Cost").fg(Color::Cyan),
220        Cell::new("Deployment Size").fg(Color::Cyan),
221    ]);
222    table.add_row(vec![num(&contract.gas, None), num(&contract.size, None)]);
223    // Add a blank row to separate deployment info from function info.
224    table.add_row(vec![Cell::new("")]);
225    table.add_row(vec![
226        Cell::new("Function Name"),
227        Cell::new("Min").fg(Color::Green),
228        Cell::new("Avg").fg(Color::Yellow),
229        Cell::new("Median").fg(Color::Yellow),
230        Cell::new("Max").fg(Color::Red),
231        Cell::new("# Calls").fg(Color::Cyan),
232    ]);
233    for (fname, sigs) in &contract.functions {
234        for (sig, gas_info) in sigs {
235            // Show function signature if overloaded else display function name.
236            let display_name = if sigs.len() == 1 { fname.clone() } else { sig.replace(':', "") };
237            table.add_row(vec![
238                Cell::new(display_name),
239                num(&gas_info.min, Some(Color::Green)),
240                num(&gas_info.mean, Some(Color::Yellow)),
241                num(&gas_info.median, Some(Color::Yellow)),
242                num(&gas_info.max, Some(Color::Red)),
243                num(&gas_info.calls, None),
244            ]);
245        }
246    }
247    table
248}
249
250#[derive(Clone, Debug, Default, Serialize, Deserialize)]
251pub struct ContractInfo {
252    pub gas: u64,
253    pub size: usize,
254    /// Function name -> Function signature -> GasInfo
255    pub functions: BTreeMap<String, BTreeMap<String, GasInfo>>,
256}
257
258#[derive(Clone, Debug, Default, Serialize, Deserialize)]
259pub struct GasInfo {
260    pub calls: u64,
261    pub min: u64,
262    pub mean: u64,
263    pub median: u64,
264    pub max: u64,
265
266    #[serde(skip)]
267    pub frames: Vec<u64>,
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273    use foundry_evm::constants::MONAD_CHEATCODE_ADDRESS;
274
275    #[test]
276    fn network_cheatcode_addresses_are_opt_in() {
277        let ethereum = GasReport::new([], [], false, []);
278        assert!(!ethereum.is_internal_address(MONAD_CHEATCODE_ADDRESS));
279
280        let monad = GasReport::new([], [], false, [MONAD_CHEATCODE_ADDRESS]);
281        assert!(monad.is_internal_address(MONAD_CHEATCODE_ADDRESS));
282    }
283}