feat(cheatcodes): stopExpectSafeMemory

crates/cheatcodes/spec/src/vm.rs

@@ -701,6 +701,10 @@ interface Vm {
     #[cheatcode(group = Testing, safety = Unsafe)]
     function expectSafeMemory(uint64 min, uint64 max) external;
 
+    /// Stops all safe memory expectation in the current subcontext.
+    #[cheatcode(group = Testing, safety = Unsafe)]
+    function stopExpectSafeMemory() external;
+
     /// Only allows memory writes to offsets [0x00, 0x60) ∪ [min, max) in the next created subcontext.
     /// If any other memory is written to, the test will fail. Can be called multiple times to add more ranges
     /// to the set.

crates/cheatcodes/src/test/expect.rs

@@ -282,6 +282,14 @@ impl Cheatcode for expectSafeMemoryCall {
     }
 }
 
+impl Cheatcode for stopExpectSafeMemoryCall {
+    fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self {} = self;
+        ccx.state.allowed_mem_writes.remove(&ccx.data.journaled_state.depth());
+        Ok(Default::default())
+    }
+}
+
 impl Cheatcode for expectSafeMemoryCallCall {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { min, max } = *self;

testdata/cheats/MemSafety.t.sol

@@ -670,6 +670,47 @@ contract MemSafetyTest is DSTest {
         }
     }
 
+    ////////////////////////////////////////////////////////////////
+    //              `stopExpectSafeMemory` cheatcode              //
+    ////////////////////////////////////////////////////////////////
+
+    /// @dev Tests that the `stopExpectSafeMemory` cheatcode works as expected.
+    function testStopExpectSafeMemory() public {
+        uint64 initPtr;
+        assembly {
+            initPtr := mload(0x40)
+        }
+
+        vm.expectSafeMemory(initPtr, initPtr + 0x20);
+        assembly {
+            // write to allowed range
+            mstore(initPtr, 0x01)
+        }
+
+        vm.stopExpectSafeMemory();
+
+        assembly {
+            // write ouside allowed range, this should be fine
+            mstore(add(initPtr, 0x20), 0x01)
+        }
+    }
+
+    /// @dev Tests that the `stopExpectSafeMemory` cheatcode does not cause violations not being noticed.
+    function testFailStopExpectSafeMemory() public {
+        uint64 initPtr;
+        assembly {
+            initPtr := mload(0x40)
+        }
+
+        vm.expectSafeMemory(initPtr, initPtr + 0x20);
+        assembly {
+            // write outside of allowed range, this should revert
+            mstore(add(initPtr, 0x20), 0x01)
+        }
+
+        vm.stopExpectSafeMemory();
+    }
+
     ////////////////////////////////////////////////////////////////
     //                          HELPERS                           //
     ////////////////////////////////////////////////////////////////