feat: mapping slot cheats (#449)

* add parent and root for stdstorage

* chore: clarify cheatcode names

* style: forge fmt

* Apply suggestions from code review

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

---------

Co-authored-by: Clouds Flowing <[email protected]>
Co-authored-by: evalir <[email protected]>

Author: 3 people

src/StdStorage.sol

@@ -168,6 +168,41 @@ library stdStorageSafe {
         return abi.decode(read(self), (int256));
     }
 
+    function parent(StdStorage storage self) internal returns (uint256, bytes32) {
+        address who = self._target;
+        uint256 field_depth = self._depth;
+        vm.startMappingRecording();
+        uint256 child = find(self) - field_depth;
+        (bool found, bytes32 key, bytes32 parent_slot) = vm.getMappingKeyAndParentOf(who, bytes32(child));
+        if (!found) {
+            revert(
+                "stdStorage read_bool(StdStorage): Cannot find parent. Make sure you give a slot and startMappingRecording() has been called."
+            );
+        }
+        return (uint256(parent_slot), key);
+    }
+
+    function root(StdStorage storage self) internal returns (uint256) {
+        address who = self._target;
+        uint256 field_depth = self._depth;
+        vm.startMappingRecording();
+        uint256 child = find(self) - field_depth;
+        bool found;
+        bytes32 root_slot;
+        bytes32 parent_slot;
+        (found,, parent_slot) = vm.getMappingKeyAndParentOf(who, bytes32(child));
+        if (!found) {
+            revert(
+                "stdStorage read_bool(StdStorage): Cannot find parent. Make sure you give a slot and startMappingRecording() has been called."
+            );
+        }
+        while (found) {
+            root_slot = parent_slot;
+            (found,, parent_slot) = vm.getMappingKeyAndParentOf(who, bytes32(root_slot));
+        }
+        return uint256(root_slot);
+    }
+
     function bytesToBytes32(bytes memory b, uint256 offset) private pure returns (bytes32) {
         bytes32 out;
 
@@ -304,6 +339,14 @@ library stdStorage {
         return stdStorageSafe.read_int(self);
     }
 
+    function parent(StdStorage storage self) internal returns (uint256, bytes32) {
+        return stdStorageSafe.parent(self);
+    }
+
+    function root(StdStorage storage self) internal returns (uint256) {
+        return stdStorageSafe.root(self);
+    }
+
     // Private function so needs to be copied over
     function bytesToBytes32(bytes memory b, uint256 offset) private pure returns (bytes32) {
         bytes32 out;

src/Vm.sol

@@ -397,12 +397,13 @@ interface VmSafe {
     function startMappingRecording() external;
     // Stops recording all map SSTOREs for later retrieval and clears the recorded data.
     function stopMappingRecording() external;
-    // Gets the length of a mapping at a given slot, for a given address.
-    function getMappingLength(address target, bytes32 slot) external returns (uint256 length);
-    // Gets the element at index idx of a mapping at a given slot, for a given address.
-    function getMappingSlotAt(address target, bytes32 slot, uint256 idx) external returns (bytes32 value);
+    // Gets the number of elements in the mapping at the given slot, for a given address.
+    function getMappingLength(address target, bytes32 mappingSlot) external returns (uint256 length);
+    // Gets the elements at index idx of the mapping at the given slot, for a given address. The
+    // index must be less than the length of the mapping (i.e. the number of keys in the mapping).
+    function getMappingSlotAt(address target, bytes32 mappingSlot, uint256 idx) external returns (bytes32 value);
     // Gets the map key and parent of a mapping at a given slot, for a given address.
-    function getMappingKeyAndParentOf(address target, bytes32 slot)
+    function getMappingKeyAndParentOf(address target, bytes32 elementSlot)
         external
         returns (bool found, bytes32 key, bytes32 parent);
     // Writes a breakpoint to jump to in the debugger

test/StdStorage.t.sol

@@ -145,6 +145,15 @@ contract StdStorageTest is Test {
         assertEq(uint256(keccak256(abi.encode(address(this), uint256(1)))), slot);
     }
 
+    function test_StorageMapAddrRoot() public {
+        (uint256 slot, bytes32 key) =
+            stdstore.target(address(test)).sig(test.map_addr.selector).with_key(address(this)).parent();
+        assertEq(address(uint160(uint256(key))), address(this));
+        assertEq(uint256(1), slot);
+        slot = stdstore.target(address(test)).sig(test.map_addr.selector).with_key(address(this)).root();
+        assertEq(uint256(1), slot);
+    }
+
     function test_StorageMapUintFound() public {
         uint256 slot = stdstore.target(address(test)).sig(test.map_uint.selector).with_key(100).find();
         assertEq(uint256(keccak256(abi.encode(100, uint256(2)))), slot);