feat: optional state diff output

script/universal/MultisigScript.sol

@@ -10,6 +10,7 @@ import {Vm} from "lib/forge-std/src/Vm.sol";
 import {IGnosisSafe, Enum} from "./IGnosisSafe.sol";
 import {Signatures} from "./Signatures.sol";
 import {Simulation} from "./Simulation.sol";
+import {StateDiff} from "./StateDiff.sol";
 
 /// @title MultisigScript
 /// @notice Script builder for Forge scripts that require signatures from Safes. Supports both non-nested
@@ -218,8 +219,12 @@ abstract contract MultisigScript is Script {
 
         (bytes[] memory datas, uint256 value) = _transactionDatas({safes: safes});
 
+        vm.startMappingRecording();
         (Vm.AccountAccess[] memory accesses, Simulation.Payload memory simPayload) =
             _simulateForSigner({safes: safes, datas: datas, value: value});
+        (StateDiff.MappingParent[] memory parents, string memory json) =
+            StateDiff.collectStateDiff(StateDiff.CollectStateDiffOpts({accesses: accesses, simPayload: simPayload}));
+        vm.stopMappingRecording();
 
         _postSign({accesses: accesses, simPayload: simPayload});
         _postCheck({accesses: accesses, simPayload: simPayload});
@@ -228,7 +233,11 @@ abstract contract MultisigScript is Script {
         for (uint256 i; i < safes.length; i++) {
             vm.store({target: safes[i], slot: SAFE_NONCE_SLOT, value: bytes32(originalNonces[i])});
         }
-        _printDataToSign({safe: safes[0], data: datas[0], value: value});
+
+        bytes memory txData = _encodeTransactionData({safe: safes[0], data: datas[0], value: value});
+        StateDiff.recordStateDiff({json: json, parents: parents, txData: txData, targetSafe: _ownerSafe()});
+
+        _printDataToSign({safe: safes[0], data: datas[0], value: value, txData: txData});
     }
 
     /// Step 1.1 (optional)
@@ -370,8 +379,7 @@ abstract contract MultisigScript is Script {
         });
     }
 
-    function _printDataToSign(address safe, bytes memory data, uint256 value) internal {
-        bytes memory txData = _encodeTransactionData({safe: safe, data: data, value: value});
+    function _printDataToSign(address safe, bytes memory data, uint256 value, bytes memory txData) internal {
         bytes32 hash = _getTransactionHash({safe: safe, data: data, value: value});
 
         emit DataToSign({data: txData});

script/universal/StateDiff.sol

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.15;
+
+import {Vm} from "lib/forge-std/src/Vm.sol";
+
+import {Simulation} from "./Simulation.sol";
+
+library StateDiff {
+    struct MappingParent {
+        bytes32 slot;
+        bytes32 parent;
+        bytes32 key;
+    }
+
+    struct CollectStateDiffOpts {
+        Vm.AccountAccess[] accesses;
+        Simulation.Payload simPayload;
+    }
+
+    /// @notice Foundry VM instance for state manipulation during simulations
+    Vm internal constant VM = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
+
+    string internal constant OBJ = "root";
+
+    function collectStateDiff(CollectStateDiffOpts memory opts)
+        internal
+        returns (MappingParent[] memory, string memory)
+    {
+        bytes memory encodedStateDiff = abi.encode(opts.accesses);
+        string memory json = VM.serializeBytes(OBJ, "stateDiff", encodedStateDiff);
+        json = VM.serializeBytes(OBJ, "overrides", abi.encode(opts.simPayload));
+
+        MappingParent[] memory parents = new MappingParent[](1);
+        // Account for the msg.sender approval override
+        parents[0] = MappingParent({
+            slot: keccak256(abi.encode(msg.sender, uint256(8))),
+            parent: bytes32(uint256(8)),
+            key: bytes32(bytes20(msg.sender))
+        });
+
+        for (uint256 i; i < opts.accesses.length; i++) {
+            for (uint256 j; j < opts.accesses[i].storageAccesses.length; j++) {
+                (bool found, bytes32 key, bytes32 parent) = VM.getMappingKeyAndParentOf(
+                    opts.accesses[i].storageAccesses[j].account, opts.accesses[i].storageAccesses[j].slot
+                );
+                if (found) {
+                    parents = _appendToParents(
+                        parents,
+                        MappingParent({slot: opts.accesses[i].storageAccesses[j].slot, parent: parent, key: key})
+                    );
+                }
+            }
+        }
+
+        return (parents, json);
+    }
+
+    function recordStateDiff(
+        string memory json,
+        MappingParent[] memory parents,
+        bytes memory txData,
+        address targetSafe
+    ) internal {
+        json = VM.serializeBytes(OBJ, "preimages", abi.encode(parents));
+        json = VM.serializeBytes(OBJ, "dataToSign", txData);
+        json = VM.serializeAddress(OBJ, "targetSafe", targetSafe);
+
+        bool shouldWrite = VM.envOr("RECORD_STATE_DIFF", false);
+        if (shouldWrite) {
+            VM.writeJson(json, "stateDiff.json");
+        }
+    }
+
+    function _appendToParents(MappingParent[] memory parents, MappingParent memory newParent)
+        private
+        pure
+        returns (MappingParent[] memory)
+    {
+        MappingParent[] memory newArr = new MappingParent[](parents.length + 1);
+        for (uint256 i; i < parents.length; i++) {
+            newArr[i] = parents[i];
+        }
+        newArr[parents.length] = newParent;
+        return newArr;
+    }
+}