Add slot derivation library (#4975)

Author: Amxx

.changeset/gorgeous-badgers-vanish.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`SlotDerivation`: Add a library of methods for derivating common storage slots.

contracts/mocks/StorageSlotMock.sol

@@ -23,6 +23,10 @@ contract StorageSlotMock {
         slot.getUint256Slot().value = value;
     }
 
+    function setInt256Slot(bytes32 slot, int256 value) public {
+        slot.getInt256Slot().value = value;
+    }
+
     function getBooleanSlot(bytes32 slot) public view returns (bool) {
         return slot.getBooleanSlot().value;
     }
@@ -39,6 +43,10 @@ contract StorageSlotMock {
         return slot.getUint256Slot().value;
     }
 
+    function getInt256Slot(bytes32 slot) public view returns (int256) {
+        return slot.getInt256Slot().value;
+    }
+
     mapping(uint256 key => string) public stringMap;
 
     function setStringSlot(bytes32 slot, string calldata value) public {

contracts/utils/Arrays.sol

@@ -4,13 +4,15 @@
 
 pragma solidity ^0.8.20;
 
+import {SlotDerivation} from "./SlotDerivation.sol";
 import {StorageSlot} from "./StorageSlot.sol";
 import {Math} from "./math/Math.sol";
 
 /**
  * @dev Collection of functions related to array types.
  */
 library Arrays {
+    using SlotDerivation for bytes32;
     using StorageSlot for bytes32;
 
     /**
@@ -379,15 +381,11 @@ library Arrays {
      */
     function unsafeAccess(address[] storage arr, uint256 pos) internal pure returns (StorageSlot.AddressSlot storage) {
         bytes32 slot;
-        // We use assembly to calculate the storage slot of the element at index `pos` of the dynamic array `arr`
-        // following https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays.
-
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(0, arr.slot)
-            slot := add(keccak256(0, 0x20), pos)
+            slot := arr.slot
         }
-        return slot.getAddressSlot();
+        return slot.deriveArray().offset(pos).getAddressSlot();
     }
 
     /**
@@ -397,15 +395,11 @@ library Arrays {
      */
     function unsafeAccess(bytes32[] storage arr, uint256 pos) internal pure returns (StorageSlot.Bytes32Slot storage) {
         bytes32 slot;
-        // We use assembly to calculate the storage slot of the element at index `pos` of the dynamic array `arr`
-        // following https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays.
-
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(0, arr.slot)
-            slot := add(keccak256(0, 0x20), pos)
+            slot := arr.slot
         }
-        return slot.getBytes32Slot();
+        return slot.deriveArray().offset(pos).getBytes32Slot();
     }
 
     /**
@@ -415,15 +409,11 @@ library Arrays {
      */
     function unsafeAccess(uint256[] storage arr, uint256 pos) internal pure returns (StorageSlot.Uint256Slot storage) {
         bytes32 slot;
-        // We use assembly to calculate the storage slot of the element at index `pos` of the dynamic array `arr`
-        // following https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays.
-
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(0, arr.slot)
-            slot := add(keccak256(0, 0x20), pos)
+            slot := arr.slot
         }
-        return slot.getUint256Slot();
+        return slot.deriveArray().offset(pos).getUint256Slot();
     }
 
     /**

contracts/utils/README.adoc

@@ -28,6 +28,7 @@ Miscellaneous contracts and libraries containing utility functions you can use t
  * {Base64}: On-chain base64 and base64URL encoding according to https://datatracker.ietf.org/doc/html/rfc4648[RFC-4648].
  * {Strings}: Common operations for strings formatting.
  * {ShortString}: Library to encode (and decode) short strings into (or from) a single bytes32 slot for optimizing costs. Short strings are limited to 31 characters.
+ * {SlotDerivation}: Methods for deriving storage slot from ERC-7201 namespaces as well as from constructions such as mapping and arrays.
  * {StorageSlot}: Methods for accessing specific storage slots formatted as common primitive types.
  * {Multicall}: Abstract contract with an utility to allow batching together multiple calls in a single transaction. Useful for allowing EOAs to perform multiple operations at once.
  * {Context}: An utility for abstracting the sender and calldata in the current execution context.
@@ -108,6 +109,8 @@ Ethereum contracts have no native concept of an interface, so applications must
 
 {{ShortStrings}}
 
+{{SlotDerivation}}
+
 {{StorageSlot}}
 
 {{Multicall}}

contracts/utils/SlotDerivation.sol

@@ -0,0 +1,161 @@
+// SPDX-License-Identifier: MIT
+// This file was procedurally generated from scripts/generate/templates/SlotDerivation.js.
+
+pragma solidity ^0.8.20;
+
+/**
+ * @dev Library for computing storage (and transient storage) locations from namespaces and deriving slots
+ * corresponding to standard patterns. The derivation method for array and mapping matches the storage layout used by
+ * the solidity language / compiler.
+ *
+ * See https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays[Solidity docs for mappings and dynamic arrays.].
+ *
+ * Example usage:
+ * ```solidity
+ * contract Example {
+ *     // Add the library methods
+ *     using StorageSlot for bytes32;
+ *     using SlotDerivation for bytes32;
+ *
+ *     // Declare a namespace
+ *     string private constant _NAMESPACE = "<namespace>" // eg. OpenZeppelin.Slot
+ *
+ *     function setValueInNamespace(uint256 key, address newValue) internal {
+ *         _NAMESPACE.erc7201Slot().deriveMapping(key).getAddressSlot().value = newValue;
+ *     }
+ *
+ *     function getValueInNamespace(uint256 key) internal view returns (address) {
+ *         return _NAMESPACE.erc7201Slot().deriveMapping(key).getAddressSlot().value;
+ *     }
+ * }
+ * ```
+ *
+ * TIP: Consider using this library along with {StorageSlot}.
+ *
+ * NOTE: This library provides a way to manipulate storage locations in a non-standard way. Tooling for checking
+ * upgrade safety will ignore the slots accessed through this library.
+ */
+library SlotDerivation {
+    /**
+     * @dev Derive an ERC-7201 slot from a string (namespace).
+     */
+    function erc7201Slot(string memory namespace) internal pure returns (bytes32 slot) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, sub(keccak256(add(namespace, 0x20), mload(namespace)), 1))
+            slot := and(keccak256(0x00, 0x20), not(0xff))
+        }
+    }
+
+    /**
+     * @dev Add an offset to a slot to get the n-th element of a structure or an array.
+     */
+    function offset(bytes32 slot, uint256 pos) internal pure returns (bytes32 result) {
+        unchecked {
+            return bytes32(uint256(slot) + pos);
+        }
+    }
+
+    /**
+     * @dev Derive the location of the first element in an array from the slot where the length is stored.
+     */
+    function deriveArray(bytes32 slot) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, slot)
+            result := keccak256(0x00, 0x20)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, address key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, key)
+            mstore(0x20, slot)
+            result := keccak256(0x00, 0x40)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, bool key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, key)
+            mstore(0x20, slot)
+            result := keccak256(0x00, 0x40)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, bytes32 key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, key)
+            mstore(0x20, slot)
+            result := keccak256(0x00, 0x40)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, uint256 key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, key)
+            mstore(0x20, slot)
+            result := keccak256(0x00, 0x40)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, int256 key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, key)
+            mstore(0x20, slot)
+            result := keccak256(0x00, 0x40)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, string memory key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let length := mload(key)
+            let begin := add(key, 0x20)
+            let end := add(begin, length)
+            let cache := mload(end)
+            mstore(end, slot)
+            result := keccak256(begin, add(length, 0x20))
+            mstore(end, cache)
+        }
+    }
+
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     */
+    function deriveMapping(bytes32 slot, bytes memory key) internal pure returns (bytes32 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let length := mload(key)
+            let begin := add(key, 0x20)
+            let end := add(begin, length)
+            let cache := mload(end)
+            mstore(end, slot)
+            result := keccak256(begin, add(length, 0x20))
+            mstore(end, cache)
+        }
+    }
+}

contracts/utils/StorageSlot.sol

@@ -15,6 +15,7 @@ pragma solidity ^0.8.20;
  * Example usage to set ERC-1967 implementation slot:
  * ```solidity
  * contract ERC1967 {
+ *     // Define the slot. Alternatively, use the SlotDerivation library to derive the slot.
  *     bytes32 internal constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
  *
  *     function _getImplementation() internal view returns (address) {
@@ -27,6 +28,8 @@ pragma solidity ^0.8.20;
  *     }
  * }
  * ```
+ *
+ * TIP: Consider using this library along with {SlotDerivation}.
  */
 library StorageSlot {
     struct AddressSlot {
@@ -45,6 +48,10 @@ library StorageSlot {
         uint256 value;
     }
 
+    struct Int256Slot {
+        int256 value;
+    }
+
     struct StringSlot {
         string value;
     }
@@ -93,6 +100,16 @@ library StorageSlot {
         }
     }
 
+    /**
+     * @dev Returns an `Int256Slot` with member `value` located at `slot`.
+     */
+    function getInt256Slot(bytes32 slot) internal pure returns (Int256Slot storage r) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            r.slot := slot
+        }
+    }
+
     /**
      * @dev Returns an `StringSlot` with member `value` located at `slot`.
      */