1use 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#[derive(Clone, Debug, Default, Serialize, Deserialize)]
20pub struct GasReport {
21 report_any: bool,
23 report_for: HashSet<String>,
25 ignore: HashSet<String>,
27 include_tests: bool,
29 #[serde(skip)]
31 extra_cheatcode_addresses: HashSet<Address>,
32 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 #[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 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 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 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 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 #[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 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 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 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 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}