add support for string and bytes keys in mapping derivation

contracts/utils/StorageSlot.sol

@@ -143,6 +143,42 @@ library StorageSlot {
         }
     }
 
+    /**
+     * @dev Derive the location of a mapping element from the key.
+     *
+     * See: https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays.
+     */
+    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.
+     *
+     * See: https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays.
+     */
+    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)
+        }
+    }
+
     /// Storage slots as structs
     struct AddressSlot {
         address value;

scripts/generate/templates/StorageSlot.js

@@ -91,6 +91,26 @@ function deriveMapping(bytes32 slot, ${type} key) internal pure returns (bytes32
 }
 `;
 
+const derive2 = ({ type }) => `\
+/**
+ * @dev Derive the location of a mapping element from the key.
+ *
+ * See: https://docs.soliditylang.org/en/v0.8.20/internals/layout_in_storage.html#mappings-and-dynamic-arrays.
+ */
+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)
+  }
+}
+`;
+
 const struct = ({ type, name }) => `\
 struct ${name}Slot {
   ${type} value;
@@ -179,7 +199,7 @@ module.exports = format(
   'library StorageSlot {',
   '/// Derivation tooling',
   tooling,
-  TYPES.filter(type => type.isValueType).flatMap(type => derive(type)), // TODO support non-value type
+  TYPES.flatMap(type => (type.isValueType ? derive(type) : derive2(type))),
   '/// Storage slots as structs',
   TYPES.flatMap(type => [struct(type), type.isValueType ? '' : getStorage(type)]),
   '/// Storage slots as udvt',

scripts/generate/templates/StorageSlot.t.js

@@ -3,7 +3,7 @@ const { capitalize } = require('../../helpers');
 const { TYPES } = require('./StorageSlot.opts');
 
 const header = `\
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.24;
 
 import {Test} from "forge-std/Test.sol";
 
@@ -40,71 +40,46 @@ function testValue${name}_2(${type} value) public {
 }
 `;
 
-const array = ({ type, name }) => `\
-${type}[] private _${type}Array;
+const array = `\
+bytes[] private _array;
 
-function testArray${name}_1(${type}[] calldata values) public {
-  bytes32 slot;
-  assembly {
-    slot := _${type}Array.slot
-  }
-
-  // set in solidity
-  _${type}Array = values;
-
-  // read using Slots
-  assertEq(slot.asUint256Slot().sload(), values.length);
-  for (uint256 i = 0; i < values.length; ++i) {
-    assertEq(slot.deriveArray().offset(i).as${name}Slot().sload(), values[i]);
-  }
-}
+function testArray(uint256 length, uint256 offset) public {
+  length = bound(length, 1, type(uint256).max);
+  offset = bound(offset, 0, length - 1);
 
-function testArray${name}_2(${type}[] calldata values) public {
-  bytes32 slot;
+  bytes32 baseSlot;
   assembly {
-    slot := _${type}Array.slot
+    baseSlot := _array.slot
   }
+  baseSlot.asUint256Slot().sstore(length);
 
-  // set using Slots
-  slot.asUint256Slot().sstore(values.length);
-  for (uint256 i = 0; i < values.length; ++i) {
-    slot.deriveArray().offset(i).as${name}Slot().sstore(values[i]);
+  bytes storage derived = _array[offset];
+  bytes32 derivedSlot;
+  assembly {
+    derivedSlot := derived.slot
   }
 
-  // read in solidity
-  assertEq(_${type}Array.length, values.length);
-  for (uint256 i = 0; i < values.length; ++i) {
-    assertEq(_${type}Array[i], values[i]);
-  }
+  assertEq(baseSlot.asUint256Slot().sload(), _array.length);
+  assertEq(baseSlot.deriveArray().offset(offset), derivedSlot);
 }
 `;
 
-const mapping = ({ type, name }) => `\
-mapping(${type} => uint256) private _${type}Mapping;
+const mapping = ({ type, name, isValueType }) => `\
+mapping(${type} => bytes) private _${type}Mapping;
 
-function testMapping${name}_1(${type} key, uint256 value) public {
-  bytes32 slot;
+function testMapping${name}(${type} ${isValueType ? '' : 'memory'} key) public {
+  bytes32 baseSlot;
   assembly {
-    slot := _${type}Mapping.slot
+    baseSlot := _${type}Mapping.slot
   }
 
-  // set in solidity
-  _${type}Mapping[key] = value;
-  // read using Slots
-  assertEq(slot.deriveMapping(key).asUint256Slot().sload(), value);
-}
-
-function testMapping${name}_2(${type} key, uint256 value) public {
-  bytes32 slot;
+  bytes storage derived = _${type}Mapping[key];
+  bytes32 derivedSlot;
   assembly {
-    slot := _${type}Mapping.slot
+    derivedSlot := derived.slot
   }
 
-  // set using Slots
-  slot.deriveMapping(key).asUint256Slot().sstore(value);
-
-  // read in solidity
-  assertEq(_${type}Mapping[key], value);
+  assertEq(baseSlot.deriveMapping(key), derivedSlot);
 }
 `;
 
@@ -115,15 +90,18 @@ module.exports = format(
   'contract StorageSlotTest is Test {',
   'using StorageSlot for *;',
   '',
-  // bool is not using a full word, solidity allocation in storage is not right aligned
+  // bool is not using a full word, solidity allocation packs such values
   TYPES.filter(type => type.isValueType && type.type !== 'bool').map(type => variable(type)),
-  // bool is not using a full word, solidity allocation in storage is not right aligned
-  TYPES.filter(type => type.isValueType && type.type !== 'bool').map(type => array(type)),
-  TYPES.filter(type => type.isValueType).flatMap(type =>
+  array,
+  TYPES.flatMap(type =>
     [].concat(
-      mapping(type),
-      (type.variants ?? []).map(variant => mapping({ type: variant, name: capitalize(variant) })),
+      type,
+      (type.variants ?? []).map(variant => ({
+        type: variant,
+        name: capitalize(variant),
+        isValueType: type.isValueType,
+      })),
     ),
-  ),
+  ).map(type => mapping(type)),
   '}',
 );