feat(cheatcodes): add resetGasMetering cheatcode (foundry-rs#8750)

* feat(cheatcodes): add resetGasMetering cheatcode

* Changes after review: nit, add test for negative gas

* Consistent gas reset if touched

crates/cheatcodes/spec/src/vm.rs

@@ -662,6 +662,10 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function resumeGasMetering() external;
 
+    /// Reset gas metering (i.e. gas usage is set to gas limit).
+    #[cheatcode(group = Evm, safety = Safe)]
+    function resetGasMetering() external;
+
     // -------- Gas Measurement --------
 
     /// Gets the gas used in the last call.

crates/cheatcodes/src/evm.rs

@@ -222,10 +222,18 @@ impl Cheatcode for resumeGasMeteringCall {
     }
 }
 
+impl Cheatcode for resetGasMeteringCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self {} = self;
+        state.gas_metering.reset();
+        Ok(Default::default())
+    }
+}
+
 impl Cheatcode for lastCallGasCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self {} = self;
-        let Some(last_call_gas) = &state.last_call_gas else {
+        let Some(last_call_gas) = &state.gas_metering.last_call_gas else {
             bail!("no external call was made yet");
         };
         Ok(last_call_gas.abi_encode())

crates/cheatcodes/src/inspector.rs

@@ -219,21 +219,36 @@ pub struct BroadcastableTransaction {
 pub struct GasMetering {
     /// True if gas metering is paused.
     pub paused: bool,
-    /// True if gas metering was resumed during the test.
-    pub resumed: bool,
+    /// True if gas metering was resumed or reseted during the test.
+    /// Used to reconcile gas when frame ends (if spent less than refunded).
+    pub touched: bool,
+    /// True if gas metering should be reset to frame limit.
+    pub reset: bool,
     /// Stores frames paused gas.
     pub paused_frames: Vec<Gas>,
+
+    /// Cache of the amount of gas used in previous call.
+    /// This is used by the `lastCallGas` cheatcode.
+    pub last_call_gas: Option<crate::Vm::Gas>,
 }
 
 impl GasMetering {
     /// Resume paused gas metering.
     pub fn resume(&mut self) {
         if self.paused {
             self.paused = false;
-            self.resumed = true;
+            self.touched = true;
         }
         self.paused_frames.clear();
     }
+
+    /// Reset gas to limit.
+    pub fn reset(&mut self) {
+        self.paused = false;
+        self.touched = true;
+        self.reset = true;
+        self.paused_frames.clear();
+    }
 }
 
 /// List of transactions that can be broadcasted.
@@ -264,7 +279,7 @@ pub struct Cheatcodes {
     /// execution block environment.
     pub block: Option<BlockEnv>,
 
-    /// The gas price
+    /// The gas price.
     ///
     /// Used in the cheatcode handler to overwrite the gas price separately from the gas price
     /// in the execution environment.
@@ -295,10 +310,6 @@ pub struct Cheatcodes {
     /// Recorded logs
     pub recorded_logs: Option<Vec<crate::Vm::Log>>,
 
-    /// Cache of the amount of gas used in previous call.
-    /// This is used by the `lastCallGas` cheatcode.
-    pub last_call_gas: Option<crate::Vm::Gas>,
-
     /// Mocked calls
     // **Note**: inner must a BTreeMap because of special `Ord` impl for `MockCallDataContext`
     pub mocked_calls: HashMap<Address, BTreeMap<MockCallDataContext, MockCallReturnData>>,
@@ -377,7 +388,6 @@ impl Cheatcodes {
             accesses: Default::default(),
             recorded_account_diffs_stack: Default::default(),
             recorded_logs: Default::default(),
-            last_call_gas: Default::default(),
             mocked_calls: Default::default(),
             expected_calls: Default::default(),
             expected_emits: Default::default(),
@@ -994,11 +1004,16 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
     fn step(&mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext<DB>) {
         self.pc = interpreter.program_counter();
 
-        // `pauseGasMetering`: reset interpreter gas.
+        // `pauseGasMetering`: pause / resume interpreter gas.
         if self.gas_metering.paused {
             self.meter_gas(interpreter);
         }
 
+        // `resetGasMetering`: reset interpreter gas.
+        if self.gas_metering.reset {
+            self.meter_gas_reset(interpreter);
+        }
+
         // `record`: record storage reads and writes.
         if self.accesses.is_some() {
             self.record_accesses(interpreter);
@@ -1026,7 +1041,7 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
             self.meter_gas_end(interpreter);
         }
 
-        if self.gas_metering.resumed {
+        if self.gas_metering.touched {
             self.meter_gas_check(interpreter);
         }
     }
@@ -1143,7 +1158,7 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
         // Record the gas usage of the call, this allows the `lastCallGas` cheatcode to
         // retrieve the gas usage of the last call.
         let gas = outcome.result.gas;
-        self.last_call_gas = Some(crate::Vm::Gas {
+        self.gas_metering.last_call_gas = Some(crate::Vm::Gas {
             gasLimit: gas.limit(),
             gasTotalUsed: gas.spent(),
             gasMemoryUsed: 0,
@@ -1406,15 +1421,22 @@ impl Cheatcodes {
         }
     }
 
+    #[cold]
+    fn meter_gas_reset(&mut self, interpreter: &mut Interpreter) {
+        interpreter.gas = Gas::new(interpreter.gas().limit());
+        self.gas_metering.reset = false;
+    }
+
     #[cold]
     fn meter_gas_check(&mut self, interpreter: &mut Interpreter) {
         if will_exit(interpreter.instruction_result) {
-            // Reconcile gas if spent is less than refunded.
-            // This can happen if gas was paused / resumed (https://github.com/foundry-rs/foundry/issues/4370).
+            // Reset gas if spent is less than refunded.
+            // This can happen if gas was paused / resumed or reset.
+            // https://github.com/foundry-rs/foundry/issues/4370
             if interpreter.gas.spent() <
                 u64::try_from(interpreter.gas.refunded()).unwrap_or_default()
             {
-                interpreter.gas = Gas::new(interpreter.gas.remaining());
+                interpreter.gas = Gas::new(interpreter.gas.limit());
             }
         }
     }

crates/forge/tests/cli/test_cmd.rs

@@ -1390,7 +1390,7 @@ contract ATest is Test {
 
     cmd.args(["test"]).with_no_redact().assert_success().stdout_eq(str![[r#"
 ...
-[PASS] test_negativeGas() (gas: 3252)
+[PASS] test_negativeGas() (gas: 0)
 ...
 "#]]);
 });
@@ -1481,3 +1481,123 @@ Traces:
 ...
 "#]]);
 });
+
+forgetest_init!(gas_metering_reset, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.clear();
+
+    prj.add_source(
+        "ATest.t.sol",
+        r#"pragma solidity 0.8.24;
+import {Vm} from "./Vm.sol";
+import {DSTest} from "./test.sol";
+contract B {
+    function a() public returns (uint256) {
+        return 100;
+    }
+}
+contract ATest is DSTest {
+    Vm vm = Vm(HEVM_ADDRESS);
+    B b;
+    uint256 a;
+
+    function testResetGas() public {
+        vm.resetGasMetering();
+    }
+
+    function testResetGas1() public {
+        vm.resetGasMetering();
+        b = new B();
+        vm.resetGasMetering();
+    }
+
+    function testResetGas2() public {
+        b = new B();
+        b = new B();
+        vm.resetGasMetering();
+    }
+
+    function testResetGas3() public {
+        vm.resetGasMetering();
+        b = new B();
+        b = new B();
+    }
+
+    function testResetGas4() public {
+        vm.resetGasMetering();
+        b = new B();
+        vm.resetGasMetering();
+        b = new B();
+    }
+
+    function testResetGas5() public {
+        vm.resetGasMetering();
+        b = new B();
+        vm.resetGasMetering();
+        b = new B();
+        vm.resetGasMetering();
+    }
+
+    function testResetGas6() public {
+        vm.resetGasMetering();
+        b = new B();
+        b = new B();
+        _reset();
+        vm.resetGasMetering();
+    }
+
+    function testResetGas7() public {
+        vm.resetGasMetering();
+        b = new B();
+        b = new B();
+        _reset();
+    }
+
+    function testResetGas8() public {
+        this.resetExternal();
+    }
+
+    function testResetGas9() public {
+        this.resetExternal();
+        vm.resetGasMetering();
+    }
+
+    function testResetNegativeGas() public {
+        a = 100;
+        vm.resetGasMetering();
+
+        delete a;
+    }
+
+    function _reset() internal {
+        vm.resetGasMetering();
+    }
+
+    function resetExternal() external {
+        b = new B();
+        b = new B();
+        vm.resetGasMetering();
+    }
+}
+     "#,
+    )
+    .unwrap();
+
+    cmd.args(["test"]).with_no_redact().assert_success().stdout_eq(str![[r#"
+...
+[PASS] testResetGas() (gas: 40)
+[PASS] testResetGas1() (gas: 40)
+[PASS] testResetGas2() (gas: 40)
+[PASS] testResetGas3() (gas: 134476)
+[PASS] testResetGas4() (gas: 56302)
+[PASS] testResetGas5() (gas: 40)
+[PASS] testResetGas6() (gas: 40)
+[PASS] testResetGas7() (gas: 49)
+[PASS] testResetGas8() (gas: 622)
+[PASS] testResetGas9() (gas: 40)
+[PASS] testResetNegativeGas() (gas: 0)
+...
+"#]]);
+});