add gas snapshot setup from #8945

.gitignore

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

crates/cheatcodes/spec/src/vm.rs

@@ -508,6 +508,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.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function snapshotGasLastCall(string calldata name) external;
+
+    /// Snapshot capture the gas usage of the last call by name in a group.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function snapshotGasLastCall(string calldata group, string calldata name) external;
+
+    /// 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.

crates/cheatcodes/src/config.rs

@@ -50,6 +50,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.
@@ -65,6 +67,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()];
@@ -93,6 +96,7 @@ impl CheatsConfig {
             labels: config.labels.clone(),
             script_wallets,
             available_artifacts,
+            running_contract,
             running_version,
             assertions_revert: config.assertions_revert,
             seed: config.fuzz.seed,
@@ -222,6 +226,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,
@@ -241,6 +246,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
     }