Add PairMap, NibbleMap and UintXMap to Bitmaps

.changeset/long-flowers-itch.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`BitMaps`: Add new map types `PairMap`, `NibbleMap`, `Uint8Map`, `Uint16Map`, `Uint32Map`, `Uint64Map`, and `Uint128Map` for gas-efficient storage of packed values with different bit widths.

contracts/utils/structs/BitMaps.sol

@@ -1,19 +1,28 @@
 // SPDX-License-Identifier: MIT
 // OpenZeppelin Contracts (last updated v5.0.0) (utils/structs/BitMaps.sol)
+// This file was procedurally generated from scripts/generate/templates/BitMaps.js.
+
 pragma solidity ^0.8.20;
 
 /**
- * @dev Library for managing uint256 to bool mapping in a compact and efficient way, provided the keys are sequential.
+ * @dev Library for managing bytes-based mappings in a compact and efficient way, provided the keys are sequential.
  * Largely inspired by Uniswap's https://github.com/Uniswap/merkle-distributor/blob/master/contracts/MerkleDistributor.sol[merkle-distributor].
  *
- * BitMaps pack 256 booleans across each bit of a single 256-bit slot of `uint256` type.
- * Hence booleans corresponding to 256 _sequential_ indices would only consume a single slot,
- * unlike the regular `bool` which would consume an entire slot for a single value.
+ * The library provides several map types that pack multiple values into single 256-bit storage slots:
+ *
+ * * `BitMap`: 256 booleans per slot (1 bit each)
+ * * `PairMap`: 128 2-bit values per slot
+ * * `NibbleMap`: 64 4-bit values per slot
+ * * `Uint8Map`: 32 8-bit values per slot
+ * * `Uint16Map`: 16 16-bit values per slot
+ * * `Uint32Map`: 8 32-bit values per slot
+ * * `Uint64Map`: 4 64-bit values per slot
+ * * `Uint128Map`: 2 128-bit values per slot
  *
- * This results in gas savings in two ways:
+ * This approach provides significant gas savings compared to using individual storage slots for each value:
  *
- * - Setting a zero value to non-zero only once every 256 times
- * - Accessing the same warm slot for every 256 _sequential_ indices
+ * * Setting a zero value to non-zero only once every N times (where N is the packing density)
+ * * Accessing the same warm slot for every N _sequential_ indices
  */
 library BitMaps {
     struct BitMap {
@@ -57,4 +66,179 @@ library BitMaps {
         uint256 mask = 1 << (index & 0xff);
         bitmap._data[bucket] &= ~mask;
     }
+
+    struct PairMap {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 2-bit value at `index` in `map`.
+     */
+    function get(PairMap storage map, uint256 index) internal view returns (uint8) {
+        uint256 bucket = index >> 7; // 128 values per bucket (256/2)
+        uint256 shift = (index & 0x7f) << 1; // i.e. (index % 128) * 2 = position * 2
+        return uint8((map._data[bucket] >> shift) & 0x03);
+    }
+
+    /**
+     * @dev Sets the 2-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 2 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(PairMap storage map, uint256 index, uint8 value) internal {
+        uint256 bucket = index >> 7; // 128 values per bucket (256/2)
+        uint256 shift = (index & 0x7f) << 1; // i.e. (index % 128) * 2 = position * 2
+        uint256 mask = 0x03 << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 2 bits
+    }
+
+    struct NibbleMap {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 4-bit value at `index` in `map`.
+     */
+    function get(NibbleMap storage map, uint256 index) internal view returns (uint8) {
+        uint256 bucket = index >> 6; // 64 values per bucket (256/4)
+        uint256 shift = (index & 0x3f) << 2; // i.e. (index % 64) * 4 = position * 4
+        return uint8((map._data[bucket] >> shift) & 0x0f);
+    }
+
+    /**
+     * @dev Sets the 4-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 4 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(NibbleMap storage map, uint256 index, uint8 value) internal {
+        uint256 bucket = index >> 6; // 64 values per bucket (256/4)
+        uint256 shift = (index & 0x3f) << 2; // i.e. (index % 64) * 4 = position * 4
+        uint256 mask = 0x0f << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 4 bits
+    }
+
+    struct Uint8Map {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 8-bit value at `index` in `map`.
+     */
+    function get(Uint8Map storage map, uint256 index) internal view returns (uint8) {
+        uint256 bucket = index >> 5; // 32 values per bucket (256/8)
+        uint256 shift = (index & 0x1f) << 3; // i.e. (index % 32) * 8 = position * 8
+        return uint8(map._data[bucket] >> shift);
+    }
+
+    /**
+     * @dev Sets the 8-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 8 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(Uint8Map storage map, uint256 index, uint8 value) internal {
+        uint256 bucket = index >> 5; // 32 values per bucket (256/8)
+        uint256 shift = (index & 0x1f) << 3; // i.e. (index % 32) * 8 = position * 8
+        uint256 mask = 0xff << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 8 bits
+    }
+
+    struct Uint16Map {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 16-bit value at `index` in `map`.
+     */
+    function get(Uint16Map storage map, uint256 index) internal view returns (uint16) {
+        uint256 bucket = index >> 4; // 16 values per bucket (256/16)
+        uint256 shift = (index & 0x0f) << 4; // i.e. (index % 16) * 16 = position * 16
+        return uint16(map._data[bucket] >> shift);
+    }
+
+    /**
+     * @dev Sets the 16-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 16 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(Uint16Map storage map, uint256 index, uint16 value) internal {
+        uint256 bucket = index >> 4; // 16 values per bucket (256/16)
+        uint256 shift = (index & 0x0f) << 4; // i.e. (index % 16) * 16 = position * 16
+        uint256 mask = 0xffff << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 16 bits
+    }
+
+    struct Uint32Map {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 32-bit value at `index` in `map`.
+     */
+    function get(Uint32Map storage map, uint256 index) internal view returns (uint32) {
+        uint256 bucket = index >> 3; // 8 values per bucket (256/32)
+        uint256 shift = (index & 0x07) << 5; // i.e. (index % 8) * 32 = position * 32
+        return uint32(map._data[bucket] >> shift);
+    }
+
+    /**
+     * @dev Sets the 32-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 32 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(Uint32Map storage map, uint256 index, uint32 value) internal {
+        uint256 bucket = index >> 3; // 8 values per bucket (256/32)
+        uint256 shift = (index & 0x07) << 5; // i.e. (index % 8) * 32 = position * 32
+        uint256 mask = 0xffffffff << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 32 bits
+    }
+
+    struct Uint64Map {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 64-bit value at `index` in `map`.
+     */
+    function get(Uint64Map storage map, uint256 index) internal view returns (uint64) {
+        uint256 bucket = index >> 2; // 4 values per bucket (256/64)
+        uint256 shift = (index & 0x03) << 6; // i.e. (index % 4) * 64 = position * 64
+        return uint64(map._data[bucket] >> shift);
+    }
+
+    /**
+     * @dev Sets the 64-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 64 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(Uint64Map storage map, uint256 index, uint64 value) internal {
+        uint256 bucket = index >> 2; // 4 values per bucket (256/64)
+        uint256 shift = (index & 0x03) << 6; // i.e. (index % 4) * 64 = position * 64
+        uint256 mask = 0xffffffffffffffff << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 64 bits
+    }
+
+    struct Uint128Map {
+        mapping(uint256 bucket => uint256) _data;
+    }
+
+    /**
+     * @dev Returns the 128-bit value at `index` in `map`.
+     */
+    function get(Uint128Map storage map, uint256 index) internal view returns (uint128) {
+        uint256 bucket = index >> 1; // 2 values per bucket (256/128)
+        uint256 shift = (index & 0x01) << 7; // i.e. (index % 2) * 128 = position * 128
+        return uint128(map._data[bucket] >> shift);
+    }
+
+    /**
+     * @dev Sets the 128-bit value at `index` in `map`.
+     *
+     * NOTE: Assumes `value` fits in 128 bits. Assembly-manipulated values may corrupt adjacent data.
+     */
+    function set(Uint128Map storage map, uint256 index, uint128 value) internal {
+        uint256 bucket = index >> 1; // 2 values per bucket (256/128)
+        uint256 shift = (index & 0x01) << 7; // i.e. (index % 2) * 128 = position * 128
+        uint256 mask = 0xffffffffffffffffffffffffffffffff << shift;
+        map._data[bucket] = (map._data[bucket] & ~mask) | (uint256(value) << shift); // set the 128 bits
+    }
 }

docs/modules/ROOT/pages/utilities.adoc

@@ -193,6 +193,81 @@ Some use cases require more powerful data structures than arrays and mappings of
 
 The `Enumerable*` structures are similar to mappings in that they store and remove elements in constant time and don't allow for repeated entries, but they also support _enumeration_, which means you can easily query all stored entries both on and off-chain.
 
+=== Efficient Storage with BitMaps
+
+The xref:api:utils.adoc#BitMaps[`BitMaps`] library provides a collection of gas-efficient data structures that pack multiple values into single storage slots. This approach can significantly reduce gas costs by minimizing storage operations, especially when dealing with sequential indices.
+
+==== Using BitMaps for Boolean Flags
+
+The most common use case is `BitMap` for storing boolean flags efficiently:
+
+[source,solidity]
+----
+using BitMaps for BitMaps.BitMap;
+
+contract VotingContract {
+    BitMaps.BitMap private _hasVoted;
+
+    function vote(uint256 proposalId) external {
+        require(!_hasVoted.get(proposalId), "Already voted");
+        _hasVoted.set(proposalId);
+        // ... voting logic
+    }
+
+    function hasVoted(uint256 proposalId) external view returns (bool) {
+        return _hasVoted.get(proposalId);
+    }
+}
+----
+
+==== Working with Small Integer Values
+
+For scenarios requiring small integer values, use the appropriate map type. For example, storing user roles that fit in 2 bits:
+
+[source,solidity]
+----
+using BitMaps for BitMaps.PairMap;
+
+contract AccessControl {
+    // 0: no access, 1: read, 2: write, 3: admin
+    BitMaps.PairMap private _userRoles;
+
+    function setUserRole(uint256 userId, uint8 role) external {
+        require(role <= 3, "Invalid role");
+        _userRoles.set(userId, role);
+    }
+
+    function getUserRole(uint256 userId) external view returns (uint8) {
+        return _userRoles.get(userId);
+    }
+}
+----
+
+==== Handling Larger Values
+
+For larger values, use the appropriate sized map. Here's an example with time delays that fit in 32 bits:
+
+[source,solidity]
+----
+using BitMaps for BitMaps.Uint32Map;
+
+contract DelaysTracker {
+    BitMaps.Uint32Map private _delays;
+
+    function recordDelay(uint256 id) external {
+        _delays.set(id, 12 hours);
+    }
+
+    function getDelay(uint256 id) external view returns (uint32) {
+        return _delays.get(id);
+    }
+}
+----
+
+IMPORTANT: Values larger than the map's bit size are truncated. For example, setting a value of 16 in a `NibbleMap` (4-bit) will result in 0. Always ensure your values fit within the allocated bits.
+
+TIP: BitMaps are most efficient when indices are sequential or clustered, as they share storage slots. Sparse indices may not provide significant gas savings over regular mappings.
+
 === Building a Merkle Tree
 
 Building an on-chain Merkle Tree allows developers to keep track of the history of roots in a decentralized manner. For these cases, the xref:api:utils.adoc#MerkleTree[`MerkleTree`] includes a predefined structure with functions to manipulate the tree (e.g. pushing values or resetting the tree).
@@ -482,4 +557,4 @@ function _setTargetAdminDelay(address target, uint32 newDelay) internal virtual
 
     emit TargetAdminDelayUpdated(target, newDelay, effect);
 }
-----
+----

scripts/generate/run.js

@@ -37,6 +37,7 @@ for (const [file, template] of Object.entries({
   'utils/structs/Checkpoints.sol': './templates/Checkpoints.js',
   'utils/structs/EnumerableSet.sol': './templates/EnumerableSet.js',
   'utils/structs/EnumerableMap.sol': './templates/EnumerableMap.js',
+  'utils/structs/BitMaps.sol': './templates/BitMaps.js',
   'utils/SlotDerivation.sol': './templates/SlotDerivation.js',
   'utils/StorageSlot.sol': './templates/StorageSlot.js',
   'utils/TransientSlot.sol': './templates/TransientSlot.js',
@@ -51,6 +52,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/structs/BitMaps.t.sol': './templates/BitMaps.t.js',
   'utils/Packing.t.sol': './templates/Packing.t.js',
   'utils/SlotDerivation.t.sol': './templates/SlotDerivation.t.js',
 })) {