Add enumerable data structures for bytes4 (#6091)

Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>

Author: ernestognw

.changeset/bright-webs-create.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`EnumerableMap`: Add support for `Bytes4ToAddressMap` types.

.changeset/thick-banks-relate.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`EnumerableSet`: Add support for `Bytes4Set` type.

contracts/utils/structs/EnumerableMap.sol

@@ -40,6 +40,7 @@ import {EnumerableSet} from "./EnumerableSet.sol";
  * - `address -> bytes32` (`AddressToBytes32Map`) since v5.1.0
  * - `bytes32 -> address` (`Bytes32ToAddressMap`) since v5.1.0
  * - `bytes -> bytes` (`BytesToBytesMap`) since v5.4.0
+ * - `bytes4 -> address` (`Bytes4ToAddressMap`) since v5.6.0
  *
  * [WARNING]
  * ====
@@ -1187,6 +1188,129 @@ library EnumerableMap {
         return result;
     }
 
+    // Bytes4ToAddressMap
+
+    struct Bytes4ToAddressMap {
+        Bytes32ToBytes32Map _inner;
+    }
+
+    /**
+     * @dev Adds a key-value pair to a map, or updates the value for an existing
+     * key. O(1).
+     *
+     * Returns true if the key was added to the map, that is if it was not
+     * already present.
+     */
+    function set(Bytes4ToAddressMap storage map, bytes4 key, address value) internal returns (bool) {
+        return set(map._inner, bytes32(key), bytes32(uint256(uint160(value))));
+    }
+
+    /**
+     * @dev Removes a value from a map. O(1).
+     *
+     * Returns true if the key was removed from the map, that is if it was present.
+     */
+    function remove(Bytes4ToAddressMap storage map, bytes4 key) internal returns (bool) {
+        return remove(map._inner, bytes32(key));
+    }
+
+    /**
+     * @dev Removes all the entries from a map. O(n).
+     *
+     * WARNING: This function has an unbounded cost that scales with map size. Developers should keep in mind that
+     * using it may render the function uncallable if the map grows to the point where clearing it consumes too much
+     * gas to fit in a block.
+     */
+    function clear(Bytes4ToAddressMap storage map) internal {
+        clear(map._inner);
+    }
+
+    /**
+     * @dev Returns true if the key is in the map. O(1).
+     */
+    function contains(Bytes4ToAddressMap storage map, bytes4 key) internal view returns (bool) {
+        return contains(map._inner, bytes32(key));
+    }
+
+    /**
+     * @dev Returns the number of elements in the map. O(1).
+     */
+    function length(Bytes4ToAddressMap storage map) internal view returns (uint256) {
+        return length(map._inner);
+    }
+
+    /**
+     * @dev Returns the element stored at position `index` in the map. O(1).
+     * Note that there are no guarantees on the ordering of values inside the
+     * array, and it may change when more values are added or removed.
+     *
+     * Requirements:
+     *
+     * - `index` must be strictly less than {length}.
+     */
+    function at(Bytes4ToAddressMap storage map, uint256 index) internal view returns (bytes4 key, address value) {
+        (bytes32 atKey, bytes32 val) = at(map._inner, index);
+        return (bytes4(atKey), address(uint160(uint256(val))));
+    }
+
+    /**
+     * @dev Tries to return the value associated with `key`. O(1).
+     * Does not revert if `key` is not in the map.
+     */
+    function tryGet(Bytes4ToAddressMap storage map, bytes4 key) internal view returns (bool exists, address value) {
+        (bool success, bytes32 val) = tryGet(map._inner, bytes32(key));
+        return (success, address(uint160(uint256(val))));
+    }
+
+    /**
+     * @dev Returns the value associated with `key`. O(1).
+     *
+     * Requirements:
+     *
+     * - `key` must be in the map.
+     */
+    function get(Bytes4ToAddressMap storage map, bytes4 key) internal view returns (address) {
+        return address(uint160(uint256(get(map._inner, bytes32(key)))));
+    }
+
+    /**
+     * @dev Returns an array containing all the keys
+     *
+     * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed
+     * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that
+     * this function has an unbounded cost, and using it as part of a state-changing function may render the function
+     * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block.
+     */
+    function keys(Bytes4ToAddressMap storage map) internal view returns (bytes4[] memory) {
+        bytes32[] memory store = keys(map._inner);
+        bytes4[] memory result;
+
+        assembly ("memory-safe") {
+            result := store
+        }
+
+        return result;
+    }
+
+    /**
+     * @dev Returns an array containing a slice of the keys
+     *
+     * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed
+     * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that
+     * this function has an unbounded cost, and using it as part of a state-changing function may render the function
+     * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block.
+     */
+    function keys(Bytes4ToAddressMap storage map, uint256 start, uint256 end) internal view returns (bytes4[] memory) {
+        bytes32[] memory store = keys(map._inner, start, end);
+        bytes4[] memory result;
+
+        assembly ("memory-safe") {
+            result := store
+        }
+
+        return result;
+    }
+
     /**
      * @dev Query for a nonexistent map key.
      */

contracts/utils/structs/EnumerableSet.sol

@@ -36,6 +36,7 @@ import {Math} from "../math/Math.sol";
  * - `uint256` (`UintSet`) since v3.3.0
  * - `string` (`StringSet`) since v5.4.0
  * - `bytes` (`BytesSet`) since v5.4.0
+ * - `bytes4` (`Bytes4Set`) since v5.6.0
  *
  * [WARNING]
  * ====
@@ -302,6 +303,108 @@ library EnumerableSet {
         return result;
     }
 
+    // Bytes4Set
+
+    struct Bytes4Set {
+        Set _inner;
+    }
+
+    /**
+     * @dev Add a value to a set. O(1).
+     *
+     * Returns true if the value was added to the set, that is if it was not
+     * already present.
+     */
+    function add(Bytes4Set storage set, bytes4 value) internal returns (bool) {
+        return _add(set._inner, bytes32(value));
+    }
+
+    /**
+     * @dev Removes a value from a set. O(1).
+     *
+     * Returns true if the value was removed from the set, that is if it was
+     * present.
+     */
+    function remove(Bytes4Set storage set, bytes4 value) internal returns (bool) {
+        return _remove(set._inner, bytes32(value));
+    }
+
+    /**
+     * @dev Removes all the values from a set. O(n).
+     *
+     * WARNING: Developers should keep in mind that this function has an unbounded cost and using it may render the
+     * function uncallable if the set grows to the point where clearing it consumes too much gas to fit in a block.
+     */
+    function clear(Bytes4Set storage set) internal {
+        _clear(set._inner);
+    }
+
+    /**
+     * @dev Returns true if the value is in the set. O(1).
+     */
+    function contains(Bytes4Set storage set, bytes4 value) internal view returns (bool) {
+        return _contains(set._inner, bytes32(value));
+    }
+
+    /**
+     * @dev Returns the number of values in the set. O(1).
+     */
+    function length(Bytes4Set storage set) internal view returns (uint256) {
+        return _length(set._inner);
+    }
+
+    /**
+     * @dev Returns the value stored at position `index` in the set. O(1).
+     *
+     * Note that there are no guarantees on the ordering of values inside the
+     * array, and it may change when more values are added or removed.
+     *
+     * Requirements:
+     *
+     * - `index` must be strictly less than {length}.
+     */
+    function at(Bytes4Set storage set, uint256 index) internal view returns (bytes4) {
+        return bytes4(_at(set._inner, index));
+    }
+
+    /**
+     * @dev Return the entire set in an array
+     *
+     * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed
+     * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that
+     * this function has an unbounded cost, and using it as part of a state-changing function may render the function
+     * uncallable if the set grows to a point where copying to memory consumes too much gas to fit in a block.
+     */
+    function values(Bytes4Set storage set) internal view returns (bytes4[] memory) {
+        bytes32[] memory store = _values(set._inner);
+        bytes4[] memory result;
+
+        assembly ("memory-safe") {
+            result := store
+        }
+
+        return result;
+    }
+
+    /**
+     * @dev Return a slice of the set in an array
+     *
+     * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed
+     * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that
+     * this function has an unbounded cost, and using it as part of a state-changing function may render the function
+     * uncallable if the set grows to a point where copying to memory consumes too much gas to fit in a block.
+     */
+    function values(Bytes4Set storage set, uint256 start, uint256 end) internal view returns (bytes4[] memory) {
+        bytes32[] memory store = _values(set._inner, start, end);
+        bytes4[] memory result;
+
+        assembly ("memory-safe") {
+            result := store
+        }
+
+        return result;
+    }
+
     // AddressSet
 
     struct AddressSet {

scripts/generate/templates/Enumerable.opts.js

@@ -24,6 +24,7 @@ const toMapTypeDescr = ({ key, value }) => ({
 
 const SET_TYPES = [
   { type: 'bytes32' },
+  { type: 'bytes4' },
   { type: 'address' },
   { type: 'uint256' },
   { type: 'string', memory: true },
@@ -38,6 +39,8 @@ const MAP_TYPES = []
     ['uint256', 'address', 'bytes32']
       .flatMap((keyType, _, array) => array.map(valueType => ({ key: { type: keyType }, value: { type: valueType } })))
       .slice(0, -1), // remove bytes32 → bytes32 (last one) that is already defined
+    // other value type maps
+    { key: { type: 'bytes4' }, value: { type: 'address' } },
     // non-value type maps
     { key: { type: 'bytes', memory: true }, value: { type: 'bytes', memory: true } },
   )

scripts/generate/templates/EnumerableMap.js

@@ -41,6 +41,7 @@ import {EnumerableSet} from "./EnumerableSet.sol";
  * - \`address -> bytes32\` (\`AddressToBytes32Map\`) since v5.1.0
  * - \`bytes32 -> address\` (\`Bytes32ToAddressMap\`) since v5.1.0
  * - \`bytes -> bytes\` (\`BytesToBytesMap\`) since v5.4.0
+ * - \`bytes4 -> address\` (\`Bytes4ToAddressMap\`) since v5.6.0
  *
  * [WARNING]
  * ====

scripts/generate/templates/EnumerableSet.js

@@ -37,6 +37,7 @@ import {Math} from "../math/Math.sol";
  * - \`uint256\` (\`UintSet\`) since v3.3.0
  * - \`string\` (\`StringSet\`) since v5.4.0
  * - \`bytes\` (\`BytesSet\`) since v5.4.0
+ * - \`bytes4\` (\`Bytes4Set\`) since v5.6.0
  *
  * [WARNING]
  * ====

scripts/generate/templates/conversion.js

@@ -2,6 +2,8 @@ function toBytes32(type, value) {
   switch (type) {
     case 'bytes32':
       return value;
+    case 'bytes4':
+      return `bytes32(${value})`;
     case 'uint256':
       return `bytes32(${value})`;
     case 'address':
@@ -15,6 +17,8 @@ function fromBytes32(type, value) {
   switch (type) {
     case 'bytes32':
       return value;
+    case 'bytes4':
+      return `bytes4(${value})`;
     case 'uint256':
       return `uint256(${value})`;
     case 'address':

test/helpers/random.js

@@ -3,6 +3,7 @@ const { ethers } = require('hardhat');
 const generators = {
   address: () => ethers.Wallet.createRandom().address,
   bytes32: () => ethers.hexlify(ethers.randomBytes(32)),
+  bytes4: () => ethers.hexlify(ethers.randomBytes(4)),
   uint256: () => ethers.toBigInt(ethers.randomBytes(32)),
   int256: () => ethers.toBigInt(ethers.randomBytes(32)) + ethers.MinInt256,
   bytes: (length = 32) => ethers.hexlify(ethers.randomBytes(length)),
@@ -11,6 +12,7 @@ const generators = {
 
 generators.address.zero = ethers.ZeroAddress;
 generators.bytes32.zero = ethers.ZeroHash;
+generators.bytes4.zero = ethers.zeroPadBytes('0x', 4);
 generators.uint256.zero = 0n;
 generators.int256.zero = 0n;
 generators.bytes.zero = '0x';