gas metering cheatcodes + modifier (#248)

Author: brockelmore

src/StdCheats.sol

@@ -431,6 +431,8 @@ abstract contract StdCheats is StdCheatsSafe {
     StdStorage private stdstore;
     Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
 
+    bool private gasMeteringOff;
+
     // Skip forward or rewind time by the specified number of seconds
     function skip(uint256 time) internal virtual {
         vm.warp(block.timestamp + time);
@@ -528,6 +530,27 @@ abstract contract StdCheats is StdCheatsSafe {
         } catch (bytes memory) {}
     }
 
+    modifier noGasMetering() {
+        vm.pauseGasMetering();
+        // To prevent turning gas monitoring back on with nested functions that use this modifier,
+        // we check if gasMetering started in the off position. If it did, we don't want to turn
+        // it back on until we exit the top level function that used the modifier
+        //
+        // i.e. funcA() noGasMetering { funcB() }, where funcB has noGasMetering as well.
+        // funcA will have `gasStartedOff` as false, funcB will have it as true,
+        // so we only turn metering back on at the end of the funcA
+        bool gasStartedOff = gasMeteringOff;
+        gasMeteringOff = true;
+
+        _;
+
+        // if gas metering was on when this modifier was called, turn it back on at the end
+        if (!gasStartedOff) {
+            gasMeteringOff = false;
+            vm.resumeGasMetering();
+        }
+    }
+
     modifier skipWhenForking() {
         if (!isFork()) {
             _;

src/Vm.sol

@@ -200,6 +200,11 @@ interface VmSafe {
     function rpcUrlStructs() external view returns (Rpc[] memory);
     // If the condition is false, discard this run's fuzz inputs and generate new ones.
     function assume(bool) external pure;
+
+    // Pauses gas metering (i.e. gas usage is not counted). Noop if already paused.
+    function pauseGasMetering() external;
+    // Resumes gas metering (i.e. gas usage is counted again). Noop if already on.
+    function resumeGasMetering() external;
 }
 
 interface Vm is VmSafe {

test/StdCheats.t.sol

@@ -219,6 +219,41 @@ contract StdCheatsTest is Test {
         receipts;
     }
 
+    function testGasMeteringModifier() public {
+        uint256 gas_start_normal = gasleft();
+        addInLoop();
+        uint256 gas_used_normal = gas_start_normal - gasleft();
+
+        uint256 gas_start_single = gasleft();
+        addInLoopNoGas();
+        uint256 gas_used_single = gas_start_single - gasleft();
+
+        uint256 gas_start_double = gasleft();
+        addInLoopNoGasNoGas();
+        uint256 gas_used_double = gas_start_double - gasleft();
+
+        emit log_named_uint("Normal gas", gas_used_normal);
+        emit log_named_uint("Single modifier gas", gas_used_single);
+        emit log_named_uint("Double modifier  gas", gas_used_double);
+        assertTrue(gas_used_double + gas_used_single < gas_used_normal);
+    }
+
+    function addInLoop() internal returns (uint256) {
+        uint256 b;
+        for (uint256 i; i < 10000; i++) {
+            b += i;
+        }
+        return b;
+    }
+
+    function addInLoopNoGas() internal noGasMetering returns (uint256) {
+        return addInLoop();
+    }
+
+    function addInLoopNoGasNoGas() internal noGasMetering returns (uint256) {
+        return addInLoopNoGas();
+    }
+
     function bytesToUint_test(bytes memory b) private pure returns (uint256) {
         uint256 number;
         for (uint256 i = 0; i < b.length; i++) {