feat: gas snapshots over arbitrary sections (#8952)

* update internal naming

* further internals

* deprecate cheats

* update Solidity tests and add dedicated test for testing deprecated cheatcodes

* clarify gas snapshots

* fix build

* final fixes

* fix build

* fix repro 6355 rename

* add gas snapshot setup from #8755

* fix build + clippy warnings

* fix cheatcodes

* account for fixed CREATE / CALL gas cost

* remove import

* add stipend

* recalculate after a - b setup

* clear call_stipend, update tests

* avoid double counting external calls

* update cheatcodes, remove debug prints

* enable assertions

* clean up tests

* clean up test names

* remove snapshot directory on `forge clean`

* do not remove all snapshots by default due to multiple test suites being able to be ran concurrently or sequentially + optimize gas snapshots check - skip if none were captured

* handle edge case where we ask to compare but file does not exist, remove snapshot directory at a top level before test suites are ran

* fix path issue when attempting removal

* Update crates/cheatcodes/src/evm.rs

Co-authored-by: Arsenii Kulikov <klkvrr@gmail.com>

* Update crates/cheatcodes/src/inspector.rs

Co-authored-by: Arsenii Kulikov <klkvrr@gmail.com>

* refactor, apply recommended changes for last_snapshot_group, last_snapshot_name

* remove gas snapshots from fuzz tests for now: this is largely due to it conflicting with the FORGE_SNAPSHOT_CHECK where it is highly likely that with different fuzzed input the gas measurement differs as well. In the future it would be an idea to capture the average gas

* fix clippy

* avoid setting to 0 unnecessarily

* use if let Some

* improve comments, clarify use of last_gas_used != 0

* fix merge conflict issue

* fix arg ordering to address group naming regression

* fix import

* move snapshot name derivation to helper

* only skip initial call w/ overhead, no special handling for call frames

* add flare test

* style nits + use helper method

---------

Co-authored-by: Arsenii Kulikov <klkvrr@gmail.com>

.gitignore

@@ -1,6 +1,7 @@
 .DS_STORE
 /target
 out/
+snapshots/
 out.json
 .idea
 .vscode

crates/cheatcodes/spec/src/vm.rs

@@ -509,6 +509,49 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Unsafe)]
     function readCallers() external returns (CallerMode callerMode, address msgSender, address txOrigin);
 
+    // ----- Arbitrary Snapshots -----
+
+    /// Snapshot capture an arbitrary numerical value by name.
+    /// The group name is derived from the contract name.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function snapshotValue(string calldata name, uint256 value) external;
+
+    /// Snapshot capture an arbitrary numerical value by name in a group.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function snapshotValue(string calldata group, string calldata name, uint256 value) external;
+
+    // -------- Gas Snapshots --------
+
+    /// Snapshot capture the gas usage of the last call by name from the callee perspective.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function snapshotGasLastCall(string calldata name) external returns (uint256 gasUsed);
+
+    /// Snapshot capture the gas usage of the last call by name in a group from the callee perspective.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function snapshotGasLastCall(string calldata group, string calldata name) external returns (uint256 gasUsed);
+
+    /// Start a snapshot capture of the current gas usage by name.
+    /// The group name is derived from the contract name.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function startSnapshotGas(string calldata name) external;
+
+    /// Start a snapshot capture of the current gas usage by name in a group.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function startSnapshotGas(string calldata group, string calldata name) external;
+
+    /// Stop the snapshot capture of the current gas by latest snapshot name, capturing the gas used since the start.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function stopSnapshotGas() external returns (uint256 gasUsed);
+
+    /// Stop the snapshot capture of the current gas usage by name, capturing the gas used since the start.
+    /// The group name is derived from the contract name.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function stopSnapshotGas(string calldata name) external returns (uint256 gasUsed);
+
+    /// Stop the snapshot capture of the current gas usage by name in a group, capturing the gas used since the start.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function stopSnapshotGas(string calldata group, string calldata name) external returns (uint256 gasUsed);
+
     // -------- State Snapshots --------
 
     /// `snapshot` is being deprecated in favor of `snapshotState`. It will be removed in future versions.
@@ -698,7 +741,7 @@ interface Vm {
 
     // -------- Gas Measurement --------
 
-    /// Gets the gas used in the last call.
+    /// Gets the gas used in the last call from the callee perspective.
     #[cheatcode(group = Evm, safety = Safe)]
     function lastCallGas() external view returns (Gas memory gas);
 

crates/cheatcodes/src/config.rs

@@ -49,6 +49,8 @@ pub struct CheatsConfig {
     /// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list.
     /// If None, no validation is performed.
     pub available_artifacts: Option<ContractsByArtifact>,
+    /// Name of the script/test contract which is currently running.
+    pub running_contract: Option<String>,
     /// Version of the script/test contract which is currently running.
     pub running_version: Option<Version>,
     /// Whether to enable legacy (non-reverting) assertions.
@@ -64,6 +66,7 @@ impl CheatsConfig {
         evm_opts: EvmOpts,
         available_artifacts: Option<ContractsByArtifact>,
         script_wallets: Option<ScriptWallets>,
+        running_contract: Option<String>,
         running_version: Option<Version>,
     ) -> Self {
         let mut allowed_paths = vec![config.root.0.clone()];
@@ -92,6 +95,7 @@ impl CheatsConfig {
             labels: config.labels.clone(),
             script_wallets,
             available_artifacts,
+            running_contract,
             running_version,
             assertions_revert: config.assertions_revert,
             seed: config.fuzz.seed,
@@ -221,6 +225,7 @@ impl Default for CheatsConfig {
             labels: Default::default(),
             script_wallets: None,
             available_artifacts: Default::default(),
+            running_contract: Default::default(),
             running_version: Default::default(),
             assertions_revert: true,
             seed: None,
@@ -240,6 +245,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
     }