fix(cheatcodes): rework gas metering pause/resume (#8736)

* fix(cheatcodes): rework gas metering pause/resume

* Changes after review: simplify paused gas recording

* Move check out of meter_gas

* Fix clippy

* Add unit test for 4523

Author: grandizzy

crates/anvil/src/eth/backend/mem/mod.rs

@@ -2676,8 +2676,8 @@ pub fn transaction_build(
     WithOtherFields::new(transaction)
 }
 
-/// Prove a storage key's existence or nonexistence in the account's storage
-/// trie.
+/// Prove a storage key's existence or nonexistence in the account's storage trie.
+///
 /// `storage_key` is the hash of the desired storage key, meaning
 /// this will only work correctly under a secure trie.
 /// `storage_key` == keccak(key)

crates/anvil/src/eth/otterscan/api.rs

@@ -39,8 +39,9 @@ pub fn mentions_address(trace: LocalizedTransactionTrace, address: Address) -> O
     }
 }
 
-/// Converts the list of traces for a transaction into the expected Otterscan format, as
-/// specified in the [`ots_traceTransaction`](https://github.com/otterscan/otterscan/blob/develop/docs/custom-jsonrpc.md#ots_tracetransaction) spec
+/// Converts the list of traces for a transaction into the expected Otterscan format.
+///
+/// Follows format specified in the [`ots_traceTransaction`](https://github.com/otterscan/otterscan/blob/develop/docs/custom-jsonrpc.md#ots_tracetransaction) spec.
 pub fn batch_build_ots_traces(traces: Vec<LocalizedTransactionTrace>) -> Vec<TraceEntry> {
     traces
         .into_iter()

crates/cheatcodes/src/evm.rs

@@ -209,17 +209,16 @@ impl Cheatcode for getRecordedLogsCall {
 impl Cheatcode for pauseGasMeteringCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self {} = self;
-        if state.gas_metering.is_none() {
-            state.gas_metering = Some(None);
-        }
+        state.pause_gas_metering = true;
         Ok(Default::default())
     }
 }
 
 impl Cheatcode for resumeGasMeteringCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self {} = self;
-        state.gas_metering = None;
+        state.pause_gas_metering = false;
+        state.paused_frame_gas = vec![];
         Ok(Default::default())
     }
 }

crates/cheatcodes/src/inspector.rs

@@ -293,18 +293,11 @@ pub struct Cheatcodes {
     /// All recorded ETH `deal`s.
     pub eth_deals: Vec<DealRecord>,
 
-    /// Holds the stored gas info for when we pause gas metering. It is an `Option<Option<..>>`
-    /// because the `call` callback in an `Inspector` doesn't get access to
-    /// the `revm::Interpreter` which holds the `revm::Gas` struct that
-    /// we need to copy. So we convert it to a `Some(None)` in `apply_cheatcode`, and once we have
-    /// the interpreter, we copy the gas struct. Then each time there is an execution of an
-    /// operation, we reset the gas.
-    pub gas_metering: Option<Option<Gas>>,
-
-    /// Holds stored gas info for when we pause gas metering, and we're entering/inside
-    /// CREATE / CREATE2 frames. This is needed to make gas meter pausing work correctly when
-    /// paused and creating new contracts.
-    pub gas_metering_create: Option<Option<Gas>>,
+    /// If true then gas metering is paused.
+    pub pause_gas_metering: bool,
+
+    /// Stores frames paused gas.
+    pub paused_frame_gas: Vec<Gas>,
 
     /// Mapping slots.
     pub mapping_slots: Option<HashMap<Address, MappingSlots>>,
@@ -353,8 +346,8 @@ impl Cheatcodes {
             context: Default::default(),
             serialized_jsons: Default::default(),
             eth_deals: Default::default(),
-            gas_metering: Default::default(),
-            gas_metering_create: Default::default(),
+            pause_gas_metering: false,
+            paused_frame_gas: Default::default(),
             mapping_slots: Default::default(),
             pc: Default::default(),
             breakpoints: Default::default(),
@@ -940,7 +933,7 @@ impl Cheatcodes {
 
 impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
     #[inline]
-    fn initialize_interp(&mut self, _interpreter: &mut Interpreter, ecx: &mut EvmContext<DB>) {
+    fn initialize_interp(&mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext<DB>) {
         // When the first interpreter is initialized we've circumvented the balance and gas checks,
         // so we apply our actual block data with the correct fees and all.
         if let Some(block) = self.block.take() {
@@ -949,14 +942,19 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
         if let Some(gas_price) = self.gas_price.take() {
             ecx.env.tx.gas_price = gas_price;
         }
+
+        // Record gas for current frame if gas metering is paused.
+        if self.pause_gas_metering {
+            self.paused_frame_gas.push(interpreter.gas);
+        }
     }
 
     #[inline]
     fn step(&mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext<DB>) {
         self.pc = interpreter.program_counter();
 
         // `pauseGasMetering`: reset interpreter gas.
-        if self.gas_metering.is_some() {
+        if self.pause_gas_metering {
             self.meter_gas(interpreter);
         }
 
@@ -1343,68 +1341,20 @@ impl<DB: DatabaseExt> InspectorExt<DB> for Cheatcodes {
 impl Cheatcodes {
     #[cold]
     fn meter_gas(&mut self, interpreter: &mut Interpreter) {
-        match &self.gas_metering {
-            None => {}
-            // Need to store gas metering.
-            Some(None) => self.gas_metering = Some(Some(interpreter.gas)),
-            Some(Some(gas)) => {
-                match interpreter.current_opcode() {
-                    opcode::CREATE | opcode::CREATE2 => {
-                        // Set we're about to enter CREATE frame to meter its gas on first opcode
-                        // inside it.
-                        self.gas_metering_create = Some(None)
-                    }
-                    opcode::STOP => {
-                        // Reset gas to value recorded when paused.
-                        interpreter.gas = *gas;
-                        self.gas_metering = None;
-                        self.gas_metering_create = None;
-                    }
-                    opcode::RETURN | opcode::SELFDESTRUCT | opcode::REVERT => {
-                        match &self.gas_metering_create {
-                            None | Some(None) => {
-                                // If we are ending current execution frame, we want to reset
-                                // interpreter gas to the value of gas spent during frame, so only
-                                // the consumed gas is erased.
-                                // ref: https://github.com/bluealloy/revm/blob/2cb991091d32330cfe085320891737186947ce5a/crates/revm/src/evm_impl.rs#L190
-                                //
-                                // It would be nice if we had access to the interpreter in
-                                // `call_end`, as we could just do this there instead.
-                                interpreter.gas = Gas::new(interpreter.gas.spent());
-
-                                // Make sure CREATE gas metering is resetted.
-                                self.gas_metering_create = None
-                            }
-                            Some(Some(gas)) => {
-                                // If this was CREATE frame, set correct gas limit. This is needed
-                                // because CREATE opcodes deduct additional gas for code storage,
-                                // and deducted amount is compared to gas limit. If we set this to
-                                // 0, the CREATE would fail with out of gas.
-                                //
-                                // If we however set gas limit to the limit of outer frame, it would
-                                // cause a panic after erasing gas cost post-create. Reason for this
-                                // is pre-create REVM records `gas_limit - (gas_limit / 64)` as gas
-                                // used, and erases costs by `remaining` gas post-create.
-                                // gas used ref: https://github.com/bluealloy/revm/blob/2cb991091d32330cfe085320891737186947ce5a/crates/revm/src/instructions/host.rs#L254-L258
-                                // post-create erase ref: https://github.com/bluealloy/revm/blob/2cb991091d32330cfe085320891737186947ce5a/crates/revm/src/instructions/host.rs#L279
-                                interpreter.gas = Gas::new(gas.limit());
-
-                                // Reset CREATE gas metering because we're about to exit its frame.
-                                self.gas_metering_create = None
-                            }
-                        }
-                    }
-                    _ => {
-                        // If just starting with CREATE opcodes, record its inner frame gas.
-                        if self.gas_metering_create == Some(None) {
-                            self.gas_metering_create = Some(Some(interpreter.gas))
-                        }
+        if let Some(paused_gas) = self.paused_frame_gas.last() {
+            // Keep gas constant if paused.
+            interpreter.gas = *paused_gas;
+        } else {
+            // Record frame paused gas.
+            self.paused_frame_gas.push(interpreter.gas);
+        }
 
-                        // Don't monitor gas changes, keep it constant.
-                        interpreter.gas = *gas;
-                    }
-                }
+        // Remove recorded gas if we exit frame.
+        match interpreter.current_opcode() {
+            opcode::STOP | opcode::RETURN | opcode::REVERT | opcode::SELFDESTRUCT => {
+                self.paused_frame_gas.pop();
             }
+            _ => {}
         }
     }
 

crates/common/src/contracts.rs

@@ -341,6 +341,8 @@ unsafe fn count_different_bytes(a: &[u8], b: &[u8]) -> usize {
     sum
 }
 
+/// Returns contract name for a given contract identifier.
+///
 /// Artifact/Contract identifier can take the following form:
 /// `<artifact file name>:<contract name>`, the `artifact file name` is the name of the json file of
 /// the contract's artifact and the contract name is the name of the solidity contract, like

crates/common/src/evm.rs

@@ -18,6 +18,7 @@ use serde::Serialize;
 pub type Breakpoints = FxHashMap<char, (Address, usize)>;
 
 /// `EvmArgs` and `EnvArgs` take the highest precedence in the Config/Figment hierarchy.
+///
 /// All vars are opt-in, their default values are expected to be set by the
 /// [`foundry_config::Config`], and are always present ([`foundry_config::Config::default`])
 ///

crates/common/src/provider/runtime_transport.rs

@@ -62,6 +62,8 @@ pub enum RuntimeTransportError {
     InvalidJwt(String),
 }
 
+/// Runtime transport that only connects on first request.
+///
 /// A runtime transport is a custom [alloy_transport::Transport] that only connects when the *first*
 /// request is made. When the first request is made, it will connect to the runtime using either an
 /// HTTP WebSocket, or IPC transport depending on the URL used.

crates/doc/src/parser/mod.rs

@@ -23,9 +23,10 @@ pub use item::{ParseItem, ParseSource};
 mod comment;
 pub use comment::{Comment, CommentTag, Comments, CommentsRef};
 
-/// The documentation parser. This type implements a [Visitor] trait. While walking the parse tree,
-/// [Parser] will collect relevant source items and corresponding doc comments. The resulting
-/// [ParseItem]s can be accessed by calling [Parser::items].
+/// The documentation parser. This type implements a [Visitor] trait.
+///
+/// While walking the parse tree, [Parser] will collect relevant source items and corresponding
+/// doc comments. The resulting [ParseItem]s can be accessed by calling [Parser::items].
 #[derive(Debug, Default)]
 pub struct Parser {
     /// Initial comments from solang parser.

crates/doc/src/preprocessor/contract_inheritance.rs

@@ -7,6 +7,7 @@ use std::{collections::HashMap, path::PathBuf};
 pub const CONTRACT_INHERITANCE_ID: PreprocessorId = PreprocessorId("contract_inheritance");
 
 /// The contract inheritance preprocessor.
+///
 /// It matches the documents with inner [`ParseSource::Contract`](crate::ParseSource) elements,
 /// iterates over their [Base](solang_parser::pt::Base)s and attempts
 /// to link them with the paths of the other contract documents.

crates/evm/core/src/backend/mod.rs

@@ -50,6 +50,7 @@ pub use snapshot::{BackendSnapshot, RevertSnapshotAction, StateSnapshot};
 type ForkDB = CacheDB<SharedBackend>;
 
 /// Represents a numeric `ForkId` valid only for the existence of the `Backend`.
+///
 /// The difference between `ForkId` and `LocalForkId` is that `ForkId` tracks pairs of `endpoint +
 /// block` which can be reused by multiple tests, whereas the `LocalForkId` is unique within a test
 pub type LocalForkId = U256;

crates/evm/evm/src/executors/invariant/shrink.rs

@@ -137,6 +137,7 @@ pub(crate) fn shrink_sequence(
 }
 
 /// Checks if the given call sequence breaks the invariant.
+///
 /// Used in shrinking phase for checking candidate sequences and in replay failures phase to test
 /// persisted failures.
 /// Returns the result of invariant check (and afterInvariant call if needed) and if sequence was

crates/evm/fuzz/src/lib.rs

@@ -299,6 +299,7 @@ impl FuzzedCases {
 }
 
 /// Fixtures to be used for fuzz tests.
+///
 /// The key represents name of the fuzzed parameter, value holds possible fuzzed values.
 /// For example, for a fixture function declared as
 /// `function fixture_sender() external returns (address[] memory senders)`

crates/evm/fuzz/src/strategies/param.rs

@@ -14,9 +14,10 @@ pub fn fuzz_param(param: &DynSolType) -> BoxedStrategy<DynSolValue> {
 }
 
 /// Given a parameter type and configured fixtures for param name, returns a strategy for generating
-/// values for that type. Fixtures can be currently generated for uint, int, address, bytes and
-/// string types and are defined for parameter name.
+/// values for that type.
 ///
+/// Fixtures can be currently generated for uint, int, address, bytes and
+/// string types and are defined for parameter name.
 /// For example, fixtures for parameter `owner` of type `address` can be defined in a function with
 /// a `function fixture_owner() public returns (address[] memory)` signature.
 ///

crates/fmt/src/visit.rs

@@ -383,6 +383,8 @@ pub trait Visitor {
     }
 }
 
+/// Visitable trait for [`solang_parser::pt`] types.
+///
 /// All [`solang_parser::pt`] types, such as [Statement], should implement the [Visitable] trait
 /// that accepts a trait [Visitor] implementation, which has various callback handles for Solidity
 /// Parse Tree nodes.

crates/forge/tests/cli/test_cmd.rs

@@ -1286,3 +1286,83 @@ contract ATest is Test {
 ...
 "#]]);
 });
+
+// see https://github.com/foundry-rs/foundry/issues/5564
+forgetest_init!(repro_5564, |prj, cmd| {
+    prj.wipe_contracts();
+
+    prj.add_test(
+        "ATest.t.sol",
+        r#"
+import {Test} from "forge-std/Test.sol";
+contract ATest is Test {
+    error MyError();
+    function testSelfMeteringRevert() public {
+        vm.pauseGasMetering();
+        vm.expectRevert(MyError.selector);
+        this.selfReverts();
+    }
+    function selfReverts() external {
+        vm.resumeGasMetering();
+        revert MyError();
+    }
+}
+   "#,
+    )
+    .unwrap();
+
+    cmd.args(["test"]).with_no_redact().assert_success().stdout_eq(str![[r#"
+...
+[PASS] testSelfMeteringRevert() (gas: 3299)
+...
+"#]]);
+});
+
+// see https://github.com/foundry-rs/foundry/issues/4523
+forgetest_init!(repro_4523, |prj, cmd| {
+    prj.wipe_contracts();
+
+    prj.add_test(
+        "ATest.t.sol",
+        r#"
+import "forge-std/Test.sol";
+contract ATest is Test {
+    mapping(uint => bytes32) map;
+
+    function test_GasMeter () public {
+        vm.pauseGasMetering();
+        for (uint i = 0; i < 10000; i++) {
+            map[i] = keccak256(abi.encode(i));
+        }
+        vm.resumeGasMetering();
+
+        for (uint i = 0; i < 10000; i++) {
+            map[i] = keccak256(abi.encode(i));
+        }
+    }
+
+    function test_GasLeft () public {
+        for (uint i = 0; i < 10000; i++) {
+            map[i] = keccak256(abi.encode(i));
+        }
+
+        uint start = gasleft();
+        for (uint i = 0; i < 10000; i++) {
+            map[i] = keccak256(abi.encode(i));
+        }
+        console2.log("Gas cost:", start - gasleft());
+    }
+}
+   "#,
+    )
+    .unwrap();
+
+    cmd.args(["test", "-vvvv"]).with_no_redact().assert_success().stdout_eq(str![[r#"
+...
+Logs:
+  Gas cost: 5754479
+...
+[PASS] test_GasMeter() (gas: 5757137)
+...
+"#]]);
+});