forge/cmd/
snapshot.rs

1use super::test;
2use crate::result::{SuiteTestResult, TestKindReport, TestOutcome};
3use alloy_primitives::{U256, map::HashMap};
4use clap::{Parser, ValueHint, builder::RangedU64ValueParser};
5use comfy_table::{
6    Cell, Color, Row, Table, modifiers::UTF8_ROUND_CORNERS, presets::ASCII_MARKDOWN,
7};
8use eyre::{Context, Result};
9use foundry_cli::utils::STATIC_FUZZ_SEED;
10use foundry_common::shell;
11use regex::Regex;
12use std::{
13    cmp::Ordering,
14    fs,
15    io::{self, BufRead},
16    path::{Path, PathBuf},
17    str::FromStr,
18    sync::LazyLock,
19};
20use yansi::Paint;
21
22/// A regex that matches a basic snapshot entry like
23/// `Test:testDeposit() (gas: 58804)`
24pub static RE_BASIC_SNAPSHOT_ENTRY: LazyLock<Regex> = LazyLock::new(|| {
25    Regex::new(r"(?P<file>(.*?)):(?P<sig>(\w+)\s*\((.*?)\))\s*\(((gas:)?\s*(?P<gas>\d+)|(runs:\s*(?P<runs>\d+),\s*μ:\s*(?P<avg>\d+),\s*~:\s*(?P<med>\d+))|(runs:\s*(?P<invruns>\d+),\s*calls:\s*(?P<calls>\d+),\s*reverts:\s*(?P<reverts>\d+)))\)").unwrap()
26});
27
28/// CLI arguments for `forge snapshot`.
29#[derive(Clone, Debug, Parser)]
30pub struct GasSnapshotArgs {
31    /// Output a diff against a pre-existing gas snapshot.
32    ///
33    /// By default, the comparison is done with .gas-snapshot.
34    #[arg(
35        conflicts_with = "snap",
36        long,
37        value_hint = ValueHint::FilePath,
38        value_name = "SNAPSHOT_FILE",
39    )]
40    diff: Option<Option<PathBuf>>,
41
42    /// Compare against a pre-existing gas snapshot, exiting with code 1 if they do not match.
43    ///
44    /// Outputs a diff if the gas snapshots do not match.
45    ///
46    /// By default, the comparison is done with .gas-snapshot.
47    #[arg(
48        conflicts_with = "diff",
49        long,
50        value_hint = ValueHint::FilePath,
51        value_name = "SNAPSHOT_FILE",
52    )]
53    check: Option<Option<PathBuf>>,
54
55    // Hidden because there is only one option
56    /// How to format the output.
57    #[arg(long, hide(true))]
58    format: Option<Format>,
59
60    /// Output file for the gas snapshot.
61    #[arg(
62        long,
63        default_value = ".gas-snapshot",
64        value_hint = ValueHint::FilePath,
65        value_name = "FILE",
66    )]
67    snap: PathBuf,
68
69    /// Tolerates gas deviations up to the specified percentage.
70    #[arg(
71        long,
72        value_parser = RangedU64ValueParser::<u32>::new().range(0..100),
73        value_name = "SNAPSHOT_THRESHOLD"
74    )]
75    tolerance: Option<u32>,
76
77    /// How to sort diff results.
78    #[arg(long, value_name = "ORDER")]
79    diff_sort: Option<DiffSortOrder>,
80
81    /// All test arguments are supported
82    #[command(flatten)]
83    pub(crate) test: test::TestArgs,
84
85    /// Additional configs for test results
86    #[command(flatten)]
87    config: GasSnapshotConfig,
88}
89
90impl GasSnapshotArgs {
91    /// Returns whether `GasSnapshotArgs` was configured with `--watch`
92    pub const fn is_watch(&self) -> bool {
93        self.test.is_watch()
94    }
95
96    /// Returns the [`watchexec::Config`] necessary to bootstrap a new watch loop.
97    pub(crate) fn watchexec_config(&self) -> Result<watchexec::Config> {
98        self.test.watchexec_config()
99    }
100
101    pub async fn run(mut self) -> Result<()> {
102        // Default to a static fuzz seed so gas snapshots are deterministic,
103        // but allow the user to override it via `--fuzz-seed`.
104        if self.test.fuzz_seed.is_none() {
105            self.test.fuzz_seed = Some(U256::from_be_bytes(STATIC_FUZZ_SEED));
106        }
107
108        let outcome = self.test.compile_and_run().await?;
109        if !shell::is_quiet()
110            && !outcome.allow_failure
111            && self.diff.is_none()
112            && self.check.is_none()
113            && outcome.failed() > 0
114        {
115            sh_eprintln!(
116                "Error: gas snapshot file \"{}\" was not written because the test run failed",
117                self.snap.display()
118            )?;
119        }
120        outcome.ensure_ok(false)?;
121        let tests = self.config.apply(outcome);
122
123        if let Some(path) = self.diff {
124            let snap = path.as_ref().unwrap_or(&self.snap);
125            let snaps = read_gas_snapshot(snap)?;
126            diff(tests, snaps, self.diff_sort.unwrap_or_default())?;
127        } else if let Some(path) = self.check {
128            let snap = path.as_ref().unwrap_or(&self.snap);
129            let snaps = read_gas_snapshot(snap)?;
130            let code = if check(tests, snaps, self.tolerance) { 0 } else { 1 };
131            std::process::exit(code)
132        } else {
133            if matches!(self.format, Some(Format::Table)) {
134                let table = build_gas_snapshot_table(&tests);
135                sh_println!("\n{}", table)?;
136            }
137            write_to_gas_snapshot_file(&tests, self.snap, self.format)?;
138        }
139        Ok(())
140    }
141}
142
143// Gas report format on stdout.
144#[derive(Clone, Debug)]
145pub enum Format {
146    Table,
147}
148
149impl FromStr for Format {
150    type Err = String;
151
152    fn from_str(s: &str) -> Result<Self, Self::Err> {
153        match s {
154            "t" | "table" => Ok(Self::Table),
155            _ => Err(format!("Unrecognized format `{s}`")),
156        }
157    }
158}
159
160/// Additional filters that can be applied on the test results
161#[derive(Clone, Debug, Default, Parser)]
162struct GasSnapshotConfig {
163    /// Sort results by gas used (ascending).
164    #[arg(long)]
165    asc: bool,
166
167    /// Sort results by gas used (descending).
168    #[arg(conflicts_with = "asc", long)]
169    desc: bool,
170
171    /// Only include tests that used more gas that the given amount.
172    #[arg(long, value_name = "MIN_GAS")]
173    min: Option<u64>,
174
175    /// Only include tests that used less gas that the given amount.
176    #[arg(long, value_name = "MAX_GAS")]
177    max: Option<u64>,
178}
179
180/// Sort order for diff output
181#[derive(Clone, Debug, Default, clap::ValueEnum)]
182enum DiffSortOrder {
183    /// Sort by percentage change (smallest to largest) - default behavior
184    #[default]
185    Percentage,
186    /// Sort by percentage change (largest to smallest)
187    PercentageDesc,
188    /// Sort by absolute gas change (smallest to largest)
189    Absolute,
190    /// Sort by absolute gas change (largest to smallest)
191    AbsoluteDesc,
192}
193
194impl GasSnapshotConfig {
195    const fn is_in_gas_range(&self, gas_used: u64) -> bool {
196        if let Some(min) = self.min
197            && gas_used < min
198        {
199            return false;
200        }
201        if let Some(max) = self.max
202            && gas_used > max
203        {
204            return false;
205        }
206        true
207    }
208
209    fn apply(&self, outcome: TestOutcome) -> Vec<SuiteTestResult> {
210        let mut tests = outcome
211            .into_tests()
212            .filter(|test| self.is_in_gas_range(test.gas_used()))
213            .flat_map(expand_invariant_snapshot_entries)
214            .collect::<Vec<_>>();
215
216        if self.asc {
217            tests.sort_by_key(|a| a.gas_used());
218        } else if self.desc {
219            tests.sort_by_key(|b| std::cmp::Reverse(b.gas_used()))
220        }
221
222        tests
223    }
224}
225
226/// Expands merged invariant campaigns into per-predicate gas snapshot rows.
227fn expand_invariant_snapshot_entries(test: SuiteTestResult) -> Vec<SuiteTestResult> {
228    if !test.result.kind.is_invariant() || test.result.invariant_predicate_results.len() <= 1 {
229        return vec![test];
230    }
231
232    test.result
233        .invariant_predicate_results
234        .iter()
235        .map(|predicate| {
236            let mut expanded = test.clone();
237            expanded.signature = format!("{}()", predicate.name);
238            expanded
239        })
240        .collect()
241}
242
243/// A general entry in a gas snapshot file
244///
245/// Has the form:
246///   `<signature>(gas:? 40181)` for normal tests
247///   `<signature>(runs: 256, μ: 40181, ~: 40181)` for fuzz tests
248///   `<signature>(runs: 256, calls: 40181, reverts: 40181)` for invariant tests
249#[derive(Clone, Debug, PartialEq, Eq)]
250pub struct GasSnapshotEntry {
251    pub contract_name: String,
252    pub signature: String,
253    pub gas_used: TestKindReport,
254}
255
256impl FromStr for GasSnapshotEntry {
257    type Err = String;
258
259    fn from_str(s: &str) -> Result<Self, Self::Err> {
260        RE_BASIC_SNAPSHOT_ENTRY
261            .captures(s)
262            .and_then(|cap| {
263                cap.name("file").and_then(|file| {
264                    cap.name("sig").and_then(|sig| {
265                        if let Some(gas) = cap.name("gas") {
266                            Some(Self {
267                                contract_name: file.as_str().to_string(),
268                                signature: sig.as_str().to_string(),
269                                gas_used: TestKindReport::Unit {
270                                    gas: gas.as_str().parse().unwrap(),
271                                },
272                            })
273                        } else if let Some(runs) = cap.name("runs") {
274                            cap.name("avg")
275                                .and_then(|avg| cap.name("med").map(|med| (runs, avg, med)))
276                                .map(|(runs, avg, med)| Self {
277                                    contract_name: file.as_str().to_string(),
278                                    signature: sig.as_str().to_string(),
279                                    gas_used: TestKindReport::Fuzz {
280                                        runs: runs.as_str().parse().unwrap(),
281                                        median_gas: med.as_str().parse().unwrap(),
282                                        mean_gas: avg.as_str().parse().unwrap(),
283                                        failed_corpus_replays: 0,
284                                    },
285                                })
286                        } else {
287                            cap.name("invruns")
288                                .and_then(|runs| {
289                                    cap.name("calls").and_then(|avg| {
290                                        cap.name("reverts").map(|med| (runs, avg, med))
291                                    })
292                                })
293                                .map(|(runs, calls, reverts)| Self {
294                                    contract_name: file.as_str().to_string(),
295                                    signature: sig.as_str().to_string(),
296                                    gas_used: TestKindReport::Invariant {
297                                        runs: runs.as_str().parse().unwrap(),
298                                        calls: calls.as_str().parse().unwrap(),
299                                        reverts: reverts.as_str().parse().unwrap(),
300                                        metrics: HashMap::default(),
301                                        failed_corpus_replays: 0,
302                                        optimization_best_value: None,
303                                    },
304                                })
305                        }
306                    })
307                })
308            })
309            .ok_or_else(|| format!("Could not extract Snapshot Entry for {s}"))
310    }
311}
312
313/// Reads a list of gas snapshot entries from a gas snapshot file.
314fn read_gas_snapshot(path: impl AsRef<Path>) -> Result<Vec<GasSnapshotEntry>> {
315    let path = path.as_ref();
316    let mut entries = Vec::new();
317    for line in io::BufReader::new(
318        fs::File::open(path)
319            .wrap_err(format!("failed to read snapshot file \"{}\"", path.display()))?,
320    )
321    .lines()
322    {
323        entries
324            .push(GasSnapshotEntry::from_str(line?.as_str()).map_err(|err| eyre::eyre!("{err}"))?);
325    }
326    Ok(entries)
327}
328
329/// Writes a series of tests to a gas snapshot file after sorting them.
330fn write_to_gas_snapshot_file(
331    tests: &[SuiteTestResult],
332    path: impl AsRef<Path>,
333    _format: Option<Format>,
334) -> Result<()> {
335    let mut reports = tests
336        .iter()
337        .map(|test| {
338            format!("{}:{} {}", test.contract_name(), test.signature, test.result.kind.report())
339        })
340        .collect::<Vec<_>>();
341
342    // sort all reports
343    reports.sort();
344
345    let content = reports.join("\n");
346    Ok(fs::write(path, content)?)
347}
348
349fn build_gas_snapshot_table(tests: &[SuiteTestResult]) -> Table {
350    let mut table = Table::new();
351    if shell::is_markdown() {
352        table.load_preset(ASCII_MARKDOWN);
353    } else {
354        table.apply_modifier(UTF8_ROUND_CORNERS);
355    }
356
357    table.set_header(vec![
358        Cell::new("Contract").fg(Color::Cyan),
359        Cell::new("Signature").fg(Color::Cyan),
360        Cell::new("Report").fg(Color::Cyan),
361    ]);
362
363    for test in tests {
364        let mut row = Row::new();
365        row.add_cell(Cell::new(test.contract_name()));
366        row.add_cell(Cell::new(&test.signature));
367        row.add_cell(Cell::new(test.result.kind.report()));
368        table.add_row(row);
369    }
370
371    table
372}
373
374/// A Gas snapshot entry diff.
375#[derive(Clone, Debug, PartialEq, Eq)]
376pub struct GasSnapshotDiff {
377    pub signature: String,
378    pub source_gas_used: TestKindReport,
379    pub target_gas_used: TestKindReport,
380}
381
382impl GasSnapshotDiff {
383    /// Returns the gas diff
384    ///
385    /// `> 0` if the source used more gas
386    /// `< 0` if the target used more gas
387    const fn gas_change(&self) -> i128 {
388        self.source_gas_used.gas() as i128 - self.target_gas_used.gas() as i128
389    }
390
391    /// Determines the percentage change
392    fn gas_diff(&self) -> f64 {
393        self.gas_change() as f64 / self.target_gas_used.gas() as f64
394    }
395}
396
397/// Compares the set of tests with an existing gas snapshot.
398///
399/// Returns true all tests match
400fn check(
401    tests: Vec<SuiteTestResult>,
402    snaps: Vec<GasSnapshotEntry>,
403    tolerance: Option<u32>,
404) -> bool {
405    let snaps = snaps
406        .into_iter()
407        .map(|s| ((s.contract_name, s.signature), s.gas_used))
408        .collect::<HashMap<_, _>>();
409    let mut has_diff = false;
410    for test in tests {
411        if let Some(target_gas) =
412            snaps.get(&(test.contract_name().to_string(), test.signature.clone())).cloned()
413        {
414            let source_gas = test.result.kind.report();
415            if !within_tolerance(source_gas.gas(), target_gas.gas(), tolerance) {
416                let _ = sh_eprintln!(
417                    "Diff in \"{}::{}\": consumed \"{}\" gas, expected \"{}\" gas ",
418                    test.contract_name(),
419                    test.signature,
420                    source_gas,
421                    target_gas
422                );
423                has_diff = true;
424            }
425        } else {
426            let _ = sh_eprintln!(
427                "No matching snapshot entry found for \"{}::{}\" in snapshot file",
428                test.contract_name(),
429                test.signature
430            );
431            has_diff = true;
432        }
433    }
434    !has_diff
435}
436
437/// Compare the set of tests with an existing gas snapshot.
438fn diff(
439    tests: Vec<SuiteTestResult>,
440    snaps: Vec<GasSnapshotEntry>,
441    sort_order: DiffSortOrder,
442) -> Result<()> {
443    let snaps = snaps
444        .into_iter()
445        .map(|s| ((s.contract_name, s.signature), s.gas_used))
446        .collect::<HashMap<_, _>>();
447    let mut diffs = Vec::with_capacity(tests.len());
448    let mut new_tests = Vec::new();
449
450    for test in tests {
451        if let Some(target_gas_used) =
452            snaps.get(&(test.contract_name().to_string(), test.signature.clone())).cloned()
453        {
454            diffs.push(GasSnapshotDiff {
455                source_gas_used: test.result.kind.report(),
456                signature: format!("{}::{}", test.contract_name(), test.signature),
457                target_gas_used,
458            });
459        } else {
460            // Track new tests
461            new_tests.push(format!("{}::{}", test.contract_name(), test.signature));
462        }
463    }
464
465    let mut increased = 0;
466    let mut decreased = 0;
467    let mut unchanged = 0;
468    let mut overall_gas_change = 0i128;
469    let mut overall_gas_used = 0i128;
470
471    // Sort based on user preference
472    match sort_order {
473        DiffSortOrder::Percentage => {
474            // Default: sort by percentage change (smallest to largest)
475            diffs.sort_by(|a, b| a.gas_diff().abs().total_cmp(&b.gas_diff().abs()));
476        }
477        DiffSortOrder::PercentageDesc => {
478            // Sort by percentage change (largest to smallest)
479            diffs.sort_by(|a, b| b.gas_diff().abs().total_cmp(&a.gas_diff().abs()));
480        }
481        DiffSortOrder::Absolute => {
482            // Sort by absolute gas change (smallest to largest)
483            diffs.sort_by_key(|d| d.gas_change().abs());
484        }
485        DiffSortOrder::AbsoluteDesc => {
486            // Sort by absolute gas change (largest to smallest)
487            diffs.sort_by_key(|d| std::cmp::Reverse(d.gas_change().abs()));
488        }
489    }
490
491    for diff in &diffs {
492        let gas_change = diff.gas_change();
493        overall_gas_change += gas_change;
494        overall_gas_used += diff.target_gas_used.gas() as i128;
495        let gas_diff = diff.gas_diff();
496
497        // Classify changes
498        if gas_change > 0 {
499            increased += 1;
500        } else if gas_change < 0 {
501            decreased += 1;
502        } else {
503            unchanged += 1;
504        }
505
506        // Display with icon and before/after values
507        let icon = if gas_change > 0 {
508            "↑".red().to_string()
509        } else if gas_change < 0 {
510            "↓".green().to_string()
511        } else {
512            "━".to_string()
513        };
514
515        sh_println!(
516            "{} {} (gas: {} → {} | {} {})",
517            icon,
518            diff.signature,
519            diff.target_gas_used.gas(),
520            diff.source_gas_used.gas(),
521            fmt_change(gas_change),
522            fmt_pct_change(gas_diff)
523        )?;
524    }
525
526    // Display new tests if any
527    if !new_tests.is_empty() {
528        sh_eprintln!("\n{}", "New tests:".yellow())?;
529        for test in new_tests {
530            sh_eprintln!("  {} {}", "+".green(), test)?;
531        }
532    }
533
534    // Summary separator
535    sh_eprintln!("\n{}", "-".repeat(80))?;
536
537    let overall_gas_diff = if overall_gas_used > 0 {
538        overall_gas_change as f64 / overall_gas_used as f64
539    } else {
540        0.0
541    };
542
543    sh_eprintln!(
544        "Total tests: {}, {} {}, {} {}, {} {}",
545        diffs.len(),
546        "↑".red().to_string(),
547        increased,
548        "↓".green().to_string(),
549        decreased,
550        "━",
551        unchanged
552    )?;
553    sh_eprintln!(
554        "Overall gas change: {} ({})",
555        fmt_change(overall_gas_change),
556        fmt_pct_change(overall_gas_diff)
557    )?;
558    Ok(())
559}
560
561fn fmt_pct_change(change: f64) -> String {
562    let change_pct = change * 100.0;
563    match change.total_cmp(&0.0) {
564        Ordering::Less => format!("{change_pct:.3}%").green().to_string(),
565        Ordering::Equal => {
566            format!("{change_pct:.3}%")
567        }
568        Ordering::Greater => format!("{change_pct:.3}%").red().to_string(),
569    }
570}
571
572fn fmt_change(change: i128) -> String {
573    match change.cmp(&0) {
574        Ordering::Less => format!("{change}").green().to_string(),
575        Ordering::Equal => change.to_string(),
576        Ordering::Greater => format!("{change}").red().to_string(),
577    }
578}
579
580/// Returns true of the difference between the gas values exceeds the tolerance
581///
582/// If `tolerance` is `None`, then this returns `true` if both gas values are equal
583fn within_tolerance(source_gas: u64, target_gas: u64, tolerance_pct: Option<u32>) -> bool {
584    if let Some(tolerance) = tolerance_pct {
585        let (hi, lo) = if source_gas > target_gas {
586            (source_gas, target_gas)
587        } else {
588            (target_gas, source_gas)
589        };
590        let diff = (1. - (lo as f64 / hi as f64)) * 100.;
591        diff < tolerance as f64
592    } else {
593        source_gas == target_gas
594    }
595}
596
597#[cfg(test)]
598mod tests {
599    use super::*;
600
601    #[test]
602    fn test_tolerance() {
603        assert!(within_tolerance(100, 105, Some(5)));
604        assert!(within_tolerance(105, 100, Some(5)));
605        assert!(!within_tolerance(100, 106, Some(5)));
606        assert!(!within_tolerance(106, 100, Some(5)));
607        assert!(within_tolerance(100, 100, None));
608    }
609
610    #[test]
611    fn can_parse_basic_gas_snapshot_entry() {
612        let s = "Test:deposit() (gas: 7222)";
613        let entry = GasSnapshotEntry::from_str(s).unwrap();
614        assert_eq!(
615            entry,
616            GasSnapshotEntry {
617                contract_name: "Test".to_string(),
618                signature: "deposit()".to_string(),
619                gas_used: TestKindReport::Unit { gas: 7222 }
620            }
621        );
622    }
623
624    #[test]
625    fn can_parse_fuzz_gas_snapshot_entry() {
626        let s = "Test:deposit() (runs: 256, μ: 100, ~:200)";
627        let entry = GasSnapshotEntry::from_str(s).unwrap();
628        assert_eq!(
629            entry,
630            GasSnapshotEntry {
631                contract_name: "Test".to_string(),
632                signature: "deposit()".to_string(),
633                gas_used: TestKindReport::Fuzz {
634                    runs: 256,
635                    median_gas: 200,
636                    mean_gas: 100,
637                    failed_corpus_replays: 0
638                }
639            }
640        );
641    }
642
643    #[test]
644    fn can_parse_invariant_gas_snapshot_entry() {
645        let s = "Test:deposit() (runs: 256, calls: 100, reverts: 200)";
646        let entry = GasSnapshotEntry::from_str(s).unwrap();
647        assert_eq!(
648            entry,
649            GasSnapshotEntry {
650                contract_name: "Test".to_string(),
651                signature: "deposit()".to_string(),
652                gas_used: TestKindReport::Invariant {
653                    runs: 256,
654                    calls: 100,
655                    reverts: 200,
656                    metrics: HashMap::default(),
657                    failed_corpus_replays: 0,
658                    optimization_best_value: None,
659                }
660            }
661        );
662    }
663
664    #[test]
665    fn can_parse_invariant_gas_snapshot_entry2() {
666        let s = "ERC20Invariants:invariantBalanceSum() (runs: 256, calls: 3840, reverts: 2388)";
667        let entry = GasSnapshotEntry::from_str(s).unwrap();
668        assert_eq!(
669            entry,
670            GasSnapshotEntry {
671                contract_name: "ERC20Invariants".to_string(),
672                signature: "invariantBalanceSum()".to_string(),
673                gas_used: TestKindReport::Invariant {
674                    runs: 256,
675                    calls: 3840,
676                    reverts: 2388,
677                    metrics: HashMap::default(),
678                    failed_corpus_replays: 0,
679                    optimization_best_value: None,
680                }
681            }
682        );
683    }
684}
forge/cmd/snapshot.rs

forge/cmd/
snapshot.rs
