Add AddressToUintMap (#3150)

* add AddressToUintMap

* Update contracts/utils/structs/EnumerableMap.sol

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>

* address comments

* lint code

* merge mocks into a single file

* add PR link to changelog entry

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>
Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>

contracts/mocks/EnumerableMapMock.sol

@@ -4,7 +4,8 @@ pragma solidity ^0.8.0;
 
 import "../utils/structs/EnumerableMap.sol";
 
-contract EnumerableMapMock {
+// UintToAddressMap
+contract UintToAddressMapMock {
     using EnumerableMap for EnumerableMap.UintToAddressMap;
 
     event OperationResult(bool result);
@@ -45,3 +46,42 @@ contract EnumerableMapMock {
         return _map.get(key, errorMessage);
     }
 }
+
+// AddressToUintMap
+contract AddressToUintMapMock {
+    using EnumerableMap for EnumerableMap.AddressToUintMap;
+
+    event OperationResult(bool result);
+
+    EnumerableMap.AddressToUintMap private _map;
+
+    function contains(address key) public view returns (bool) {
+        return _map.contains(key);
+    }
+
+    function set(address key, uint256 value) public {
+        bool result = _map.set(key, value);
+        emit OperationResult(result);
+    }
+
+    function remove(address 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 (address key, uint256 value) {
+        return _map.at(index);
+    }
+
+    function tryGet(address key) public view returns (bool, uint256) {
+        return _map.tryGet(key);
+    }
+
+    function get(address key) public view returns (uint256) {
+        return _map.get(key);
+    }
+}

contracts/utils/structs/EnumerableMap.sol

@@ -26,8 +26,10 @@ import "./EnumerableSet.sol";
  * }
  * ```
  *
- * As of v3.0.0, only maps of type `uint256 -> address` (`UintToAddressMap`) are
- * supported.
+ * The following map types are supported:
+ *
+ * - `uint256 -> address` (`UintToAddressMap`) since v3.0.0
+ * - `address -> uint256` (`AddressToUintMap`) since v4.6.0
  */
 library EnumerableMap {
     using EnumerableSet for EnumerableSet.Bytes32Set;
@@ -237,4 +239,84 @@ library EnumerableMap {
     ) internal view returns (address) {
         return address(uint160(uint256(_get(map._inner, bytes32(key), errorMessage))));
     }
+
+    // AddressToUintMap
+
+    struct AddressToUintMap {
+        Map _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(
+        AddressToUintMap storage map,
+        address key,
+        uint256 value
+    ) internal returns (bool) {
+        return _set(map._inner, bytes32(uint256(uint160(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(AddressToUintMap storage map, address key) internal returns (bool) {
+        return _remove(map._inner, bytes32(uint256(uint160(key))));
+    }
+
+    /**
+     * @dev Returns true if the key is in the map. O(1).
+     */
+    function contains(AddressToUintMap storage map, address key) internal view returns (bool) {
+        return _contains(map._inner, bytes32(uint256(uint160(key))));
+    }
+
+    /**
+     * @dev Returns the number of elements in the map. O(1).
+     */
+    function length(AddressToUintMap 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(AddressToUintMap storage map, uint256 index) internal view returns (address, uint256) {
+        (bytes32 key, bytes32 value) = _at(map._inner, index);
+        return (address(uint160(uint256(key))), uint256(value));
+    }
+
+    /**
+     * @dev Tries to returns the value associated with `key`.  O(1).
+     * Does not revert if `key` is not in the map.
+     *
+     * _Available since v3.4._
+     */
+    function tryGet(AddressToUintMap storage map, address key) internal view returns (bool, uint256) {
+        (bool success, bytes32 value) = _tryGet(map._inner, bytes32(uint256(uint160(key))));
+        return (success, uint256(value));
+    }
+
+    /**
+     * @dev Returns the value associated with `key`.  O(1).
+     *
+     * Requirements:
+     *
+     * - `key` must be in the map.
+     */
+    function get(AddressToUintMap storage map, address key) internal view returns (uint256) {
+        return uint256(_get(map._inner, bytes32(uint256(uint160(key)))));
+    }
 }