Add non-value types in EnumerableSet and EnumerableMap (#5658)

Co-authored-by: Hadrien Croubois <[email protected]>

Author: ernestognw

.changeset/long-hornets-mate.md

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

.changeset/pink-dolls-shop.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`EnumerableSet`: Add support for `StringSet` and `BytesSet` types.

contracts/utils/structs/EnumerableMap.sol

@@ -39,6 +39,7 @@ import {EnumerableSet} from "./EnumerableSet.sol";
  * - `address -> address` (`AddressToAddressMap`) since v5.1.0
  * - `address -> bytes32` (`AddressToBytes32Map`) since v5.1.0
  * - `bytes32 -> address` (`Bytes32ToAddressMap`) since v5.1.0
+ * - `bytes -> bytes` (`BytesToBytesMap`) since v5.4.0
  *
  * [WARNING]
  * ====
@@ -51,7 +52,7 @@ import {EnumerableSet} from "./EnumerableSet.sol";
  * ====
  */
 library EnumerableMap {
-    using EnumerableSet for EnumerableSet.Bytes32Set;
+    using EnumerableSet for *;
 
     // To implement this library for multiple types with as little code repetition as possible, we write it in
     // terms of a generic Map type with bytes32 keys and values. The Map implementation uses private functions,
@@ -997,4 +998,122 @@ library EnumerableMap {
 
         return result;
     }
+
+    /**
+     * @dev Query for a nonexistent map key.
+     */
+    error EnumerableMapNonexistentBytesKey(bytes key);
+
+    struct BytesToBytesMap {
+        // Storage of keys
+        EnumerableSet.BytesSet _keys;
+        mapping(bytes key => bytes) _values;
+    }
+
+    /**
+     * @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(BytesToBytesMap storage map, bytes memory key, bytes memory value) internal returns (bool) {
+        map._values[key] = value;
+        return map._keys.add(key);
+    }
+
+    /**
+     * @dev Removes a key-value pair from a map. O(1).
+     *
+     * Returns true if the key was removed from the map, that is if it was present.
+     */
+    function remove(BytesToBytesMap storage map, bytes memory key) internal returns (bool) {
+        delete map._values[key];
+        return map._keys.remove(key);
+    }
+
+    /**
+     * @dev Removes all the entries from a map. 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 map grows to the point where clearing it consumes too much gas to fit in a block.
+     */
+    function clear(BytesToBytesMap storage map) internal {
+        uint256 len = length(map);
+        for (uint256 i = 0; i < len; ++i) {
+            delete map._values[map._keys.at(i)];
+        }
+        map._keys.clear();
+    }
+
+    /**
+     * @dev Returns true if the key is in the map. O(1).
+     */
+    function contains(BytesToBytesMap storage map, bytes memory key) internal view returns (bool) {
+        return map._keys.contains(key);
+    }
+
+    /**
+     * @dev Returns the number of key-value pairs in the map. O(1).
+     */
+    function length(BytesToBytesMap storage map) internal view returns (uint256) {
+        return map._keys.length();
+    }
+
+    /**
+     * @dev Returns the key-value pair stored at position `index` in the map. O(1).
+     *
+     * Note that there are no guarantees on the ordering of entries inside the
+     * array, and it may change when more entries are added or removed.
+     *
+     * Requirements:
+     *
+     * - `index` must be strictly less than {length}.
+     */
+    function at(
+        BytesToBytesMap storage map,
+        uint256 index
+    ) internal view returns (bytes memory key, bytes memory value) {
+        key = map._keys.at(index);
+        value = map._values[key];
+    }
+
+    /**
+     * @dev Tries to returns the value associated with `key`. O(1).
+     * Does not revert if `key` is not in the map.
+     */
+    function tryGet(
+        BytesToBytesMap storage map,
+        bytes memory key
+    ) internal view returns (bool exists, bytes memory value) {
+        value = map._values[key];
+        exists = bytes(value).length != 0 || contains(map, key);
+    }
+
+    /**
+     * @dev Returns the value associated with `key`. O(1).
+     *
+     * Requirements:
+     *
+     * - `key` must be in the map.
+     */
+    function get(BytesToBytesMap storage map, bytes memory key) internal view returns (bytes memory value) {
+        bool exists;
+        (exists, value) = tryGet(map, key);
+        if (!exists) {
+            revert EnumerableMapNonexistentBytesKey(key);
+        }
+    }
+
+    /**
+     * @dev Return the 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(BytesToBytesMap storage map) internal view returns (bytes[] memory) {
+        return map._keys.values();
+    }
 }