Add bytes32 to uint enumerable map (#3416)

CHANGELOG.md

@@ -8,6 +8,7 @@
  * `ERC20FlashMint`: Add customizable flash fee receiver. ([#3327](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3327))
  * `Strings`: add a new overloaded function `toHexString` that converts an `address` with fixed length of 20 bytes to its not checksummed ASCII `string` hexadecimal representation. ([#3403](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3403))
  * `EnumerableMap`: add new `UintToUintMap` map type. ([#3338](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3338))
+ * `EnumerableMap`: add new `Bytes32ToUintMap` map type. ([#3416](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3416))
  * `SafeCast`: add support for many more types, using procedural code generation. ([#3245](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3245))
 
 ## 4.6.0 (2022-04-26)

contracts/mocks/EnumerableMapMock.sol

@@ -174,3 +174,46 @@ contract UintToUintMapMock {
         return _map.get(key, errorMessage);
     }
 }
+
+// Bytes32ToUintMap
+contract Bytes32ToUintMapMock {
+    using EnumerableMap for EnumerableMap.Bytes32ToUintMap;
+
+    event OperationResult(bool result);
+
+    EnumerableMap.Bytes32ToUintMap private _map;
+
+    function contains(bytes32 key) public view returns (bool) {
+        return _map.contains(key);
+    }
+
+    function set(bytes32 key, uint256 value) public {
+        bool result = _map.set(key, value);
+        emit OperationResult(result);
+    }
+
+    function remove(bytes32 key) public {
+        bool result = _map.remove(key);
+        emit OperationResult(result);
+    }
+
+    function length() public view returns (uint256) {
+        return _map.length();
+    }
+
+    function at(uint256 index) public view returns (bytes32 key, uint256 value) {
+        return _map.at(index);
+    }
+
+    function tryGet(bytes32 key) public view returns (bool, uint256) {
+        return _map.tryGet(key);
+    }
+
+    function get(bytes32 key) public view returns (uint256) {
+        return _map.get(key);
+    }
+
+    function getWithMessage(bytes32 key, string calldata errorMessage) public view returns (uint256) {
+        return _map.get(key, errorMessage);
+    }
+}

contracts/utils/structs/EnumerableMap.sol

@@ -32,6 +32,7 @@ import "./EnumerableSet.sol";
  * - `address -> uint256` (`AddressToUintMap`) since v4.6.0
  * - `bytes32 -> bytes32` (`Bytes32ToBytes32`) since v4.6.0
  * - `uint256 -> uint256` (`UintToUintMap`) since v4.7.0
+ * - `bytes32 -> uint256` (`Bytes32ToUintMap`) since v4.7.0
  *
  * [WARNING]
  * ====
@@ -433,4 +434,96 @@ library EnumerableMap {
     ) internal view returns (uint256) {
         return uint256(get(map._inner, bytes32(uint256(uint160(key))), errorMessage));
     }
+
+    // Bytes32ToUintMap
+
+    struct Bytes32ToUintMap {
+        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(
+        Bytes32ToUintMap storage map,
+        bytes32 key,
+        uint256 value
+    ) internal returns (bool) {
+        return set(map._inner, key, bytes32(value));
+    }
+
+    /**
+     * @dev Removes a value from a set. O(1).
+     *
+     * Returns true if the key was removed from the map, that is if it was present.
+     */
+    function remove(Bytes32ToUintMap storage map, bytes32 key) internal returns (bool) {
+        return remove(map._inner, key);
+    }
+
+    /**
+     * @dev Returns true if the key is in the map. O(1).
+     */
+    function contains(Bytes32ToUintMap storage map, bytes32 key) internal view returns (bool) {
+        return contains(map._inner, key);
+    }
+
+    /**
+     * @dev Returns the number of elements in the map. O(1).
+     */
+    function length(Bytes32ToUintMap storage map) internal view returns (uint256) {
+        return length(map._inner);
+    }
+
+    /**
+     * @dev Returns the element 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(Bytes32ToUintMap storage map, uint256 index) internal view returns (bytes32, uint256) {
+        (bytes32 key, bytes32 value) = at(map._inner, index);
+        return (key, uint256(value));
+    }
+
+    /**
+     * @dev Tries to returns the value associated with `key`.  O(1).
+     * Does not revert if `key` is not in the map.
+     */
+    function tryGet(Bytes32ToUintMap storage map, bytes32 key) internal view returns (bool, uint256) {
+        (bool success, bytes32 value) = tryGet(map._inner, key);
+        return (success, uint256(value));
+    }
+
+    /**
+     * @dev Returns the value associated with `key`.  O(1).
+     *
+     * Requirements:
+     *
+     * - `key` must be in the map.
+     */
+    function get(Bytes32ToUintMap storage map, bytes32 key) internal view returns (uint256) {
+        return uint256(get(map._inner, key));
+    }
+
+    /**
+     * @dev Same as {get}, with a custom error message when `key` is not in the map.
+     *
+     * CAUTION: This function is deprecated because it requires allocating memory for the error
+     * message unnecessarily. For custom revert reasons use {tryGet}.
+     */
+    function get(
+        Bytes32ToUintMap storage map,
+        bytes32 key,
+        string memory errorMessage
+    ) internal view returns (uint256) {
+        return uint256(get(map._inner, key, errorMessage));
+    }
 }

test/utils/structs/EnumerableMap.test.js

@@ -4,6 +4,7 @@ const AddressToUintMapMock = artifacts.require('AddressToUintMapMock');
 const UintToAddressMapMock = artifacts.require('UintToAddressMapMock');
 const Bytes32ToBytes32MapMock = artifacts.require('Bytes32ToBytes32MapMock');
 const UintToUintMapMock = artifacts.require('UintToUintMapMock');
+const Bytes32ToUintMapMock = artifacts.require('Bytes32ToUintMapMock');
 
 const { shouldBehaveLikeMap } = require('./EnumerableMap.behavior');
 
@@ -25,8 +26,8 @@ contract('EnumerableMap', function (accounts) {
     });
 
     shouldBehaveLikeMap(
-      [accountA, accountB, accountC],
-      [keyA, keyB, keyC],
+      [ accountA, accountB, accountC ],
+      [ keyA, keyB, keyC ],
       new BN('0'),
     );
   });
@@ -38,8 +39,8 @@ contract('EnumerableMap', function (accounts) {
     });
 
     shouldBehaveLikeMap(
-      [keyA, keyB, keyC],
-      [accountA, accountB, accountC],
+      [ keyA, keyB, keyC ],
+      [ accountA, accountB, accountC ],
       constants.ZERO_ADDRESS,
     );
   });
@@ -51,8 +52,8 @@ contract('EnumerableMap', function (accounts) {
     });
 
     shouldBehaveLikeMap(
-      [keyA, keyB, keyC].map(k => ('0x' + k.toString(16)).padEnd(66, '0')),
-      [bytesA, bytesB, bytesC],
+      [ keyA, keyB, keyC ].map(k => '0x' + k.toString(16).padEnd(64, '0')),
+      [ bytesA, bytesB, bytesC ],
       constants.ZERO_BYTES32,
     );
   });
@@ -69,4 +70,17 @@ contract('EnumerableMap', function (accounts) {
       new BN('0'),
     );
   });
+
+  // Bytes32ToUintMap
+  describe('Bytes32ToUintMap', function () {
+    beforeEach(async function () {
+      this.map = await Bytes32ToUintMapMock.new();
+    });
+
+    shouldBehaveLikeMap(
+      [ bytesA, bytesB, bytesC ],
+      [ keyA, keyB, keyC ],
+      new BN('0'),
+    );
+  });
 });