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

forge/cmd/
snapshot.rs
