Add slot derivation library (#4975)

scripts/generate/run.js

@@ -36,6 +36,7 @@ for (const [file, template] of Object.entries({
   'utils/structs/EnumerableSet.sol': './templates/EnumerableSet.js',
   'utils/structs/EnumerableMap.sol': './templates/EnumerableMap.js',
   'utils/structs/Checkpoints.sol': './templates/Checkpoints.js',
+  'utils/SlotDerivation.sol': './templates/SlotDerivation.js',
   'utils/StorageSlot.sol': './templates/StorageSlot.js',
   'utils/Arrays.sol': './templates/Arrays.js',
 })) {
@@ -45,6 +46,7 @@ for (const [file, template] of Object.entries({
 // Tests
 for (const [file, template] of Object.entries({
   'utils/structs/Checkpoints.t.sol': './templates/Checkpoints.t.js',
+  'utils/SlotDerivation.t.sol': './templates/SlotDerivation.t.js',
 })) {
   generateFromTemplate(file, template, './test/');
 }

scripts/generate/templates/Arrays.js

@@ -5,6 +5,7 @@ const { TYPES } = require('./Arrays.opts');
 const header = `\
 pragma solidity ^0.8.20;
 
+import {SlotDerivation} from "./SlotDerivation.sol";
 import {StorageSlot} from "./StorageSlot.sol";
 import {Math} from "./math/Math.sol";
 
@@ -327,16 +328,12 @@ const unsafeAccessStorage = type => `
 function unsafeAccess(${type}[] storage arr, uint256 pos) internal pure returns (StorageSlot.${capitalize(
   type,
 )}Slot 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)
-   }
-   return slot.get${capitalize(type)}Slot();
+    bytes32 slot;
+    /// @solidity memory-safe-assembly
+    assembly {
+        slot := arr.slot
+    }
+    return slot.deriveArray().offset(pos).get${capitalize(type)}Slot();
 }`;
 
 const unsafeAccessMemory = type => `
@@ -368,6 +365,7 @@ function unsafeSetLength(${type}[] storage array, uint256 len) internal {
 module.exports = format(
   header.trimEnd(),
   'library Arrays {',
+  'using SlotDerivation for bytes32;',
   'using StorageSlot for bytes32;',
   // sorting, comparator, helpers and internal
   sort('bytes32'),

scripts/generate/templates/Slot.opts.js

@@ -0,0 +1,13 @@
+const { capitalize } = require('../../helpers');
+
+const TYPES = [
+  { type: 'address', isValueType: true },
+  { type: 'bool', isValueType: true, name: 'Boolean' },
+  { type: 'bytes32', isValueType: true, variants: ['bytes4'] },
+  { type: 'uint256', isValueType: true, variants: ['uint32'] },
+  { type: 'int256', isValueType: true, variants: ['int32'] },
+  { type: 'string', isValueType: false },
+  { type: 'bytes', isValueType: false },
+].map(type => Object.assign(type, { name: type.name ?? capitalize(type.type) }));
+
+module.exports = { TYPES };

scripts/generate/templates/SlotDerivation.js

@@ -0,0 +1,116 @@
+const format = require('../format-lines');
+const { TYPES } = require('./Slot.opts');
+
+const header = `\
+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.
+ */
+`;
+
+const namespace = `\
+/**
+ * @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))
+  }
+}
+`;
+
+const array = `\
+/**
+ * @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)
+  }
+}
+`;
+
+const mapping = ({ type }) => `\
+/**
+ * @dev Derive the location of a mapping element from the key.
+ */
+function deriveMapping(bytes32 slot, ${type} key) internal pure returns (bytes32 result) {
+  /// @solidity memory-safe-assembly
+  assembly {
+      mstore(0x00, key)
+      mstore(0x20, slot)
+      result := keccak256(0x00, 0x40)
+  }
+}
+`;
+
+const mapping2 = ({ type }) => `\
+/**
+ * @dev Derive the location of a mapping element from the key.
+ */
+function deriveMapping(bytes32 slot, ${type} 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)
+  }
+}
+`;
+
+// GENERATE
+module.exports = format(
+  header.trimEnd(),
+  'library SlotDerivation {',
+  namespace,
+  array,
+  TYPES.map(type => (type.isValueType ? mapping(type) : mapping2(type))),
+  '}',
+);

scripts/generate/templates/SlotDerivation.t.js

@@ -0,0 +1,73 @@
+const format = require('../format-lines');
+const { capitalize } = require('../../helpers');
+const { TYPES } = require('./Slot.opts');
+
+const header = `\
+pragma solidity ^0.8.20;
+
+import {Test} from "forge-std/Test.sol";
+
+import {SlotDerivation} from "@openzeppelin/contracts/utils/SlotDerivation.sol";
+`;
+
+const array = `\
+bytes[] private _array;
+
+function testDeriveArray(uint256 length, uint256 offset) public {
+  length = bound(length, 1, type(uint256).max);
+  offset = bound(offset, 0, length - 1);
+
+  bytes32 baseSlot;
+  assembly {
+    baseSlot := _array.slot
+    sstore(baseSlot, length) // store length so solidity access does not revert
+  }
+
+  bytes storage derived = _array[offset];
+  bytes32 derivedSlot;
+  assembly {
+    derivedSlot := derived.slot
+  }
+
+  assertEq(baseSlot.deriveArray().offset(offset), derivedSlot);
+}
+`;
+
+const mapping = ({ type, name, isValueType }) => `\
+mapping(${type} => bytes) private _${type}Mapping;
+
+function testDeriveMapping${name}(${type} ${isValueType ? '' : 'memory'} key) public {
+  bytes32 baseSlot;
+  assembly {
+    baseSlot := _${type}Mapping.slot
+  }
+
+  bytes storage derived = _${type}Mapping[key];
+  bytes32 derivedSlot;
+  assembly {
+    derivedSlot := derived.slot
+  }
+
+  assertEq(baseSlot.deriveMapping(key), derivedSlot);
+}
+`;
+
+// GENERATE
+module.exports = format(
+  header.trimEnd(),
+  'contract SlotDerivationTest is Test {',
+  'using SlotDerivation for bytes32;',
+  '',
+  array,
+  TYPES.flatMap(type =>
+    [].concat(
+      type,
+      (type.variants ?? []).map(variant => ({
+        type: variant,
+        name: capitalize(variant),
+        isValueType: type.isValueType,
+      })),
+    ),
+  ).map(type => mapping(type)),
+  '}',
+);

scripts/generate/templates/StorageSlot.js

@@ -1,14 +1,5 @@
 const format = require('../format-lines');
-const { capitalize } = require('../../helpers');
-
-const TYPES = [
-  { type: 'address', isValueType: true },
-  { type: 'bool', isValueType: true, name: 'Boolean' },
-  { type: 'bytes32', isValueType: true },
-  { type: 'uint256', isValueType: true },
-  { type: 'string', isValueType: false },
-  { type: 'bytes', isValueType: false },
-].map(type => Object.assign(type, { struct: (type.name ?? capitalize(type.type)) + 'Slot' }));
+const { TYPES } = require('./Slot.opts');
 
 const header = `\
 pragma solidity ^0.8.20;
@@ -24,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) {
@@ -36,20 +28,22 @@ pragma solidity ^0.8.20;
  *     }
  * }
  * \`\`\`
+ * 
+ * TIP: Consider using this library along with {SlotDerivation}.
  */
 `;
 
-const struct = type => `\
-struct ${type.struct} {
-  ${type.type} value;
+const struct = ({ type, name }) => `\
+struct ${name}Slot {
+  ${type} value;
 }
 `;
 
-const get = type => `\
+const get = ({ name }) => `\
 /**
- * @dev Returns an \`${type.struct}\` with member \`value\` located at \`slot\`.
+ * @dev Returns an \`${name}Slot\` with member \`value\` located at \`slot\`.
  */
-function get${type.struct}(bytes32 slot) internal pure returns (${type.struct} storage r) {
+function get${name}Slot(bytes32 slot) internal pure returns (${name}Slot storage r) {
   /// @solidity memory-safe-assembly
   assembly {
       r.slot := slot
@@ -57,11 +51,11 @@ function get${type.struct}(bytes32 slot) internal pure returns (${type.struct} s
 }
 `;
 
-const getStorage = type => `\
+const getStorage = ({ type, name }) => `\
 /**
- * @dev Returns an \`${type.struct}\` representation of the ${type.type} storage pointer \`store\`.
+ * @dev Returns an \`${name}Slot\` representation of the ${type} storage pointer \`store\`.
  */
-function get${type.struct}(${type.type} storage store) internal pure returns (${type.struct} storage r) {
+function get${name}Slot(${type} storage store) internal pure returns (${name}Slot storage r) {
   /// @solidity memory-safe-assembly
   assembly {
       r.slot := store.slot
@@ -73,6 +67,7 @@ function get${type.struct}(${type.type} storage store) internal pure returns (${
 module.exports = format(
   header.trimEnd(),
   'library StorageSlot {',
-  [...TYPES.map(struct), ...TYPES.flatMap(type => [get(type), type.isValueType ? '' : getStorage(type)])],
+  TYPES.map(type => struct(type)),
+  TYPES.flatMap(type => [get(type), type.isValueType ? '' : getStorage(type)]),
   '}',
 );