feat(forge): add vm.stopRecord (#10370)

* add initial implementation of StopRecord cheatcode

* modify command to be stopRecordAndReturnAccesses instead

* improve documentation and lint fixes

* lint changes in integration tests

* implement resetRecord and stopRecord

* clippy fix

* minor formatting changes

* chore: cleanup, rm reset

---------

Co-authored-by: DaniPopes <[email protected]>

Author: tushar994

crates/cheatcodes/assets/cheatcodes.json

@@ -399,10 +399,15 @@ interface Vm {
 
     // -------- Record Storage --------
 
-    /// Records all storage reads and writes.
+    /// Records all storage reads and writes. Use `accesses` to get the recorded data.
+    /// Subsequent calls to `record` will clear the previous data.
     #[cheatcode(group = Evm, safety = Safe)]
     function record() external;
 
+    /// Stops recording storage reads and writes.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function stopRecord() external;
+
     /// Gets all accessed reads and write slot from a `vm.record` session, for a given address.
     #[cheatcode(group = Evm, safety = Safe)]
     function accesses(address target) external returns (bytes32[] memory readSlots, bytes32[] memory writeSlots);

crates/cheatcodes/spec/src/vm.rs

@@ -56,6 +56,12 @@ impl RecordAccess {
         self.record_read(target, slot);
         self.writes.entry(target).or_default().push(slot);
     }
+
+    /// Clears the recorded reads and writes.
+    pub fn clear(&mut self) {
+        // Also frees memory.
+        *self = Default::default();
+    }
 }
 
 /// Records the `snapshotGas*` cheatcodes.
@@ -280,24 +286,26 @@ impl Cheatcode for dumpStateCall {
 impl Cheatcode for recordCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self {} = self;
-        state.accesses = Some(Default::default());
+        state.recording_accesses = true;
+        state.accesses.clear();
+        Ok(Default::default())
+    }
+}
+
+impl Cheatcode for stopRecordCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        state.recording_accesses = false;
         Ok(Default::default())
     }
 }
 
 impl Cheatcode for accessesCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { target } = *self;
-        let result = state
-            .accesses
-            .as_mut()
-            .map(|accesses| {
-                (
-                    &accesses.reads.entry(target).or_default()[..],
-                    &accesses.writes.entry(target).or_default()[..],
-                )
-            })
-            .unwrap_or_default();
+        let result = (
+            state.accesses.reads.entry(target).or_default().as_slice(),
+            state.accesses.writes.entry(target).or_default().as_slice(),
+        );
         Ok(result.abi_encode_params())
     }
 }

crates/cheatcodes/src/evm.rs

@@ -420,7 +420,10 @@ pub struct Cheatcodes {
     pub fork_revert_diagnostic: Option<RevertDiagnostic>,
 
     /// Recorded storage reads and writes
-    pub accesses: Option<RecordAccess>,
+    pub accesses: RecordAccess,
+
+    /// Whether storage access recording is currently active
+    pub recording_accesses: bool,
 
     /// Recorded account accesses (calls, creates) organized by relative call depth, where the
     /// topmost vector corresponds to accesses at the depth at which account access recording
@@ -538,6 +541,7 @@ impl Cheatcodes {
             assume_no_revert: Default::default(),
             fork_revert_diagnostic: Default::default(),
             accesses: Default::default(),
+            recording_accesses: Default::default(),
             recorded_account_diffs_stack: Default::default(),
             recorded_logs: Default::default(),
             record_debug_steps_info: Default::default(),
@@ -1332,7 +1336,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
         }
 
         // `record`: record storage reads and writes.
-        if self.accesses.is_some() {
+        if self.recording_accesses {
             self.record_accesses(interpreter);
         }
 
@@ -1945,7 +1949,7 @@ impl Cheatcodes {
     /// Records storage slots reads and writes.
     #[cold]
     fn record_accesses(&mut self, interpreter: &mut Interpreter) {
-        let Some(access) = &mut self.accesses else { return };
+        let access = &mut self.accesses;
         match interpreter.current_opcode() {
             op::SLOAD => {
                 let key = try_or_return!(interpreter.stack().peek(0));

crates/cheatcodes/src/inspector.rs

@@ -53,4 +53,44 @@ contract RecordTest is DSTest {
         assertEq(innerWrites.length, 1, "number of nested writes is incorrect");
         assertEq(innerWrites[0], bytes32(uint256(2)), "key for nested write is incorrect");
     }
+
+    function testStopRecordAccess() public {
+        RecordAccess target = new RecordAccess();
+
+        // Start recording
+        vm.record();
+        NestedRecordAccess inner = target.record();
+
+        // Verify Records
+        (bytes32[] memory reads, bytes32[] memory writes) = vm.accesses(address(target));
+
+        assertEq(reads.length, 2, "number of reads is incorrect");
+        assertEq(reads[0], bytes32(uint256(1)), "key for read 0 is incorrect");
+        assertEq(reads[1], bytes32(uint256(1)), "key for read 1 is incorrect");
+
+        assertEq(writes.length, 1, "number of writes is incorrect");
+        assertEq(writes[0], bytes32(uint256(1)), "key for write is incorrect");
+
+        vm.stopRecord();
+        inner = target.record();
+
+        // Verify that there are no new Records
+        (reads, writes) = vm.accesses(address(target));
+
+        assertEq(reads.length, 2, "number of reads is incorrect");
+        assertEq(reads[0], bytes32(uint256(1)), "key for read 0 is incorrect");
+        assertEq(reads[1], bytes32(uint256(1)), "key for read 1 is incorrect");
+
+        assertEq(writes.length, 1, "number of writes is incorrect");
+        assertEq(writes[0], bytes32(uint256(1)), "key for write is incorrect");
+
+        vm.record();
+        vm.stopRecord();
+
+        // verify reset all records
+        (reads, writes) = vm.accesses(address(target));
+
+        assertEq(reads.length, 0, "number of reads is incorrect");
+        assertEq(writes.length, 0, "number of writes is incorrect");
+    }
 }