feat: add way to disable emitting of gas snapshots to disk (#9710)

* feat: gas snapshot emit config

* review change comment

---------

Co-authored-by: turbocrime <[email protected]>

Author: turbocrime

crates/config/src/lib.rs

@@ -202,6 +202,8 @@ pub struct Config {
     pub snapshots: PathBuf,
     /// whether to check for differences against previously stored gas snapshots
     pub gas_snapshot_check: bool,
+    /// whether to emit gas snapshots to disk
+    pub gas_snapshot_emit: bool,
     /// where the broadcast logs are stored
     pub broadcast: PathBuf,
     /// additional solc allow paths for `--allow-paths`
@@ -2321,6 +2323,7 @@ impl Default for Config {
             broadcast: "broadcast".into(),
             snapshots: "snapshots".into(),
             gas_snapshot_check: false,
+            gas_snapshot_emit: true,
             allow_paths: vec![],
             include_paths: vec![],
             force: false,

crates/forge/bin/cmd/test/mod.rs

@@ -122,6 +122,10 @@ pub struct TestArgs {
     #[arg(long, env = "FORGE_SNAPSHOT_CHECK")]
     gas_snapshot_check: Option<bool>,
 
+    /// Enable/disable recording of gas snapshot results.
+    #[arg(long, env = "FORGE_SNAPSHOT_EMIT")]
+    gas_snapshot_emit: Option<bool>,
+
     /// Exit with code 0 even if a test fails.
     #[arg(long, env = "FORGE_ALLOW_FAILURE")]
     allow_failure: bool,
@@ -732,17 +736,28 @@ impl TestArgs {
                     }
                 }
 
-                // Create `snapshots` directory if it doesn't exist.
-                fs::create_dir_all(&config.snapshots)?;
-
-                // Write gas snapshots to disk per group.
-                gas_snapshots.clone().into_iter().for_each(|(group, snapshots)| {
-                    fs::write_pretty_json_file(
-                        &config.snapshots.join(format!("{group}.json")),
-                        &snapshots,
-                    )
-                    .expect("Failed to write gas snapshots to disk");
-                });
+                // By default `gas_snapshot_emit` is set to `true` in the config.
+                //
+                // The user can either:
+                // - Set `FORGE_SNAPSHOT_EMIT=false` in the environment.
+                // - Pass `--gas-snapshot-emit=false` as a CLI argument.
+                // - Set `gas_snapshot_emit = false` in the config.
+                //
+                // If the user passes `--gas-snapshot-emit=<bool>` then it will override the config
+                // and the environment variable, enabling the check if `true` is passed.
+                if self.gas_snapshot_emit.unwrap_or(config.gas_snapshot_emit) {
+                    // Create `snapshots` directory if it doesn't exist.
+                    fs::create_dir_all(&config.snapshots)?;
+
+                    // Write gas snapshots to disk per group.
+                    gas_snapshots.clone().into_iter().for_each(|(group, snapshots)| {
+                        fs::write_pretty_json_file(
+                            &config.snapshots.join(format!("{group}.json")),
+                            &snapshots,
+                        )
+                        .expect("Failed to write gas snapshots to disk");
+                    });
+                }
             }
 
             // Print suite summary.

crates/forge/tests/cli/config.rs

@@ -41,6 +41,7 @@ forgetest!(can_extract_config_values, |prj, cmd| {
         cache_path: "test-cache".into(),
         snapshots: "snapshots".into(),
         gas_snapshot_check: false,
+        gas_snapshot_emit: true,
         broadcast: "broadcast".into(),
         force: true,
         evm_version: EvmVersion::Byzantium,
@@ -966,6 +967,7 @@ cache = true
 cache_path = "cache"
 snapshots = "snapshots"
 gas_snapshot_check = false
+gas_snapshot_emit = true
 broadcast = "broadcast"
 allow_paths = []
 include_paths = []
@@ -1122,6 +1124,7 @@ exclude = []
   "cache_path": "cache",
   "snapshots": "snapshots",
   "gas_snapshot_check": false,
+  "gas_snapshot_emit": true,
   "broadcast": "broadcast",
   "allow_paths": [],
   "include_paths": [],
@@ -1535,3 +1538,105 @@ Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
 ...
 "#]]);
 });
+
+forgetest_init!(test_gas_snapshot_emit_config, |prj, cmd| {
+    // Default settings: gas_snapshot_emit enabled.
+    cmd.forge_fuse().args(["config"]).assert_success().stdout_eq(str![[r#"
+...
+gas_snapshot_emit = true
+...
+"#]]);
+
+    prj.insert_ds_test();
+
+    prj.add_source(
+        "GasSnapshotEmitTest.sol",
+        r#"
+import "./test.sol";
+
+interface Vm {
+    function startSnapshotGas(string memory name) external;
+    function stopSnapshotGas() external returns (uint256);
+}
+
+contract GasSnapshotEmitTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testSnapshotGasSection() public {
+        vm.startSnapshotGas("testSection");
+        int n = 1;
+        vm.stopSnapshotGas();
+    }
+}
+    "#,
+    )
+    .unwrap();
+
+    // Assert that gas_snapshot_emit is enabled by default.
+    cmd.forge_fuse().args(["test"]).assert_success();
+    // Assert that snapshots were emitted to disk.
+    assert!(prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Remove the snapshot file.
+    fs::remove_file(prj.root().join("snapshots/GasSnapshotEmitTest.json")).unwrap();
+
+    // Test that `--gas-snapshot-emit=false` flag can be used to disable writing snapshots.
+    cmd.forge_fuse().args(["test", "--gas-snapshot-emit=false"]).assert_success();
+    // Assert that snapshots were not emitted to disk.
+    assert!(!prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Test that environment variable `FORGE_SNAPSHOT_EMIT` can be used to disable writing
+    // snapshots.
+    cmd.forge_fuse();
+    cmd.env("FORGE_SNAPSHOT_EMIT", "false");
+    cmd.args(["test"]).assert_success();
+    // Assert that snapshots were not emitted to disk.
+    assert!(!prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Test that `--gas-snapshot-emit=true` flag can be used to enable writing snapshots, even when
+    // `FORGE_SNAPSHOT_EMIT` is set to false.
+    cmd.forge_fuse();
+    cmd.env("FORGE_SNAPSHOT_EMIT", "false");
+    cmd.args(["test", "--gas-snapshot-emit=true"]).assert_success();
+    // Assert that snapshots were emitted to disk.
+    assert!(prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Remove the snapshot file.
+    fs::remove_file(prj.root().join("snapshots/GasSnapshotEmitTest.json")).unwrap();
+
+    // Disable gas_snapshot_emit in the config file.
+    prj.update_config(|config| config.gas_snapshot_emit = false);
+    cmd.forge_fuse().args(["config"]).assert_success();
+
+    // Test that snapshots are not emitted to disk, when disabled by config.
+    cmd.forge_fuse().args(["test"]).assert_success();
+    // Assert that snapshots were not emitted to disk.
+    assert!(!prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Test that `--gas-snapshot-emit=true` flag can be used to enable writing snapshots, when
+    // disabled by config.
+    cmd.forge_fuse();
+    cmd.args(["test", "--gas-snapshot-emit=true"]).assert_success();
+    // Assert that snapshots were emitted to disk.
+    assert!(prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Remove the snapshot file.
+    fs::remove_file(prj.root().join("snapshots/GasSnapshotEmitTest.json")).unwrap();
+
+    // Test that environment variable `FORGE_SNAPSHOT_EMIT` can be used to enable writing snapshots.
+    cmd.forge_fuse();
+    cmd.env("FORGE_SNAPSHOT_EMIT", "true");
+    cmd.args(["test"]).assert_success();
+    // Assert that snapshots were emitted to disk.
+    assert!(prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+
+    // Remove the snapshot file.
+    fs::remove_file(prj.root().join("snapshots/GasSnapshotEmitTest.json")).unwrap();
+
+    // Test that `--gas-snapshot-emit=false` flag can be used to disable writing snapshots,
+    // even when `FORGE_SNAPSHOT_EMIT` is set to true.
+    cmd.forge_fuse().args(["test", "--gas-snapshot-emit=false"]).assert_success();
+
+    // Assert that snapshots were not emitted to disk.
+    assert!(!prj.root().join("snapshots/GasSnapshotEmitTest.json").exists());
+});