Add AddressToUintMap

.prettierrc

@@ -1,8 +1,11 @@
 {
+  "singleQuote": true,
+  "trailingComma": "all",
   "overrides": [
     {
       "files": "*.sol",
       "options": {
+        "singleQuote": false,
         "printWidth": 120,
         "explicitTypes": "always"
       }

contracts/mocks/EnumerableMap/AddressToUintMapMock.sol

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../../utils/structs/EnumerableMap.sol";
+
+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);
+    }
+
+    function getWithMessage(address key, string calldata errorMessage) public view returns (uint256) {
+        return _map.get(key, errorMessage);
+    }
+}

contracts/mocks/EnumerableMapMock.sol → contracts/mocks/EnumerableMap/UintToAddressMapMock.sol

@@ -2,9 +2,9 @@
 
 pragma solidity ^0.8.0;
 
-import "../utils/structs/EnumerableMap.sol";
+import "../../utils/structs/EnumerableMap.sol";
 
-contract EnumerableMapMock {
+contract UintToAddressMapMock {
     using EnumerableMap for EnumerableMap.UintToAddressMap;
 
     event OperationResult(bool result);

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.
+ * As of v4.5.0, the following map types are supported:
+ *
+ * - `uint256 -> address` (`UintToAddressMap`)
+ * - `address -> uint256` (`AddressToUintMap`)
  */
 library EnumerableMap {
     using EnumerableSet for EnumerableSet.Bytes32Set;
@@ -237,4 +239,106 @@ 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, addressToBytes32(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, addressToBytes32(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, addressToBytes32(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 (bytes32ToAddress(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, addressToBytes32(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, addressToBytes32(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(
+        AddressToUintMap storage map,
+        address key,
+        string memory errorMessage
+    ) internal view returns (uint256) {
+        return uint256(_get(map._inner, addressToBytes32(key), errorMessage));
+    }
+}
+
+function addressToBytes32(address _address) pure returns (bytes32) {
+    return bytes32(uint256(uint160(_address)));
+}
+
+function bytes32ToAddress(bytes32 _bytes32) pure returns (address) {
+    return address(uint160(uint256(_bytes32)));
 }

test/utils/structs/EnumerableMap/AddressToUintMap.test.js

@@ -0,0 +1,159 @@
+const { BN, expectEvent, expectRevert } = require('@openzeppelin/test-helpers');
+const { expect } = require('chai');
+const { expectMembersMatch } = require('./helpers');
+
+const AddressToUintMapMock = artifacts.require('AddressToUintMapMock');
+
+contract('AddressToUintMap', function (accounts) {
+  const [accountA, accountB, accountC] = accounts;
+
+  const valueA = new BN('7891');
+  const valueB = new BN('451');
+  const valueC = new BN('9592328');
+
+  beforeEach(async function () {
+    this.map = await AddressToUintMapMock.new();
+  });
+
+  it('starts empty', async function () {
+    expect(await this.map.contains(accountA)).to.equal(false);
+
+    await expectMembersMatch(this.map, [], []);
+  });
+
+  describe('set', function () {
+    it('adds a key', async function () {
+      const receipt = await this.map.set(accountA, valueA);
+      expectEvent(receipt, 'OperationResult', { result: true });
+
+      await expectMembersMatch(this.map, [accountA], [valueA]);
+    });
+
+    it('adds several keys', async function () {
+      await this.map.set(accountA, valueA);
+      await this.map.set(accountB, valueB);
+
+      await expectMembersMatch(this.map, [accountA, accountB], [valueA, valueB]);
+
+      expect(await this.map.contains(accountC)).to.equal(false);
+    });
+
+    it('returns false when adding keys already in the set', async function () {
+      await this.map.set(accountA, valueA);
+
+      const receipt = await this.map.set(accountA, valueA);
+      expectEvent(receipt, 'OperationResult', { result: false });
+
+      await expectMembersMatch(this.map, [accountA], [valueA]);
+    });
+
+    it('updates values for keys already in the set', async function () {
+      await this.map.set(accountA, valueA);
+      await this.map.set(accountA, valueB);
+
+      await expectMembersMatch(this.map, [accountA], [valueB]);
+    });
+  });
+
+  describe('remove', function () {
+    it('removes added keys', async function () {
+      await this.map.set(accountA, valueA);
+
+      const receipt = await this.map.remove(accountA);
+      expectEvent(receipt, 'OperationResult', { result: true });
+
+      expect(await this.map.contains(accountA)).to.equal(false);
+      await expectMembersMatch(this.map, [], []);
+    });
+
+    it('returns false when removing keys not in the set', async function () {
+      const receipt = await this.map.remove(accountA);
+      expectEvent(receipt, 'OperationResult', { result: false });
+
+      expect(await this.map.contains(accountA)).to.equal(false);
+    });
+
+    it('adds and removes multiple keys', async function () {
+      // []
+
+      await this.map.set(accountA, valueA);
+      await this.map.set(accountC, valueC);
+
+      // [A, C]
+
+      await this.map.remove(accountA);
+      await this.map.remove(accountB);
+
+      // [C]
+
+      await this.map.set(accountB, valueB);
+
+      // [C, B]
+
+      await this.map.set(accountA, valueA);
+      await this.map.remove(accountC);
+
+      // [A, B]
+
+      await this.map.set(accountA, valueA);
+      await this.map.set(accountB, valueB);
+
+      // [A, B]
+
+      await this.map.set(accountC, valueC);
+      await this.map.remove(accountA);
+
+      // [B, C]
+
+      await this.map.set(accountA, valueA);
+      await this.map.remove(accountB);
+
+      // [A, C]
+
+      await expectMembersMatch(this.map, [accountA, accountC], [valueA, valueC]);
+
+      expect(await this.map.contains(accountB)).to.equal(false);
+    });
+  });
+
+  describe('read', function () {
+    beforeEach(async function () {
+      await this.map.set(accountA, valueA);
+    });
+
+    describe('get', function () {
+      it('existing value', async function () {
+        expect(await this.map.get(accountA)).to.bignumber.equal(valueA);
+      });
+      it('missing value', async function () {
+        await expectRevert(this.map.get(accountB), 'EnumerableMap: nonexistent key');
+      });
+    });
+
+    describe('get with message', function () {
+      it('existing value', async function () {
+        expect(await this.map.getWithMessage(accountA, 'custom error string')).to.bignumber.equal(valueA);
+      });
+      it('missing value', async function () {
+        await expectRevert(this.map.getWithMessage(accountB, 'custom error string'), 'custom error string');
+      });
+    });
+
+    describe('tryGet', function () {
+      const stringifyTryGetValue = ({ 0: result, 1: value }) => ({ 0: result, 1: value.toString() });
+
+      it('existing value', async function () {
+        const actual = stringifyTryGetValue(await this.map.tryGet(accountA));
+        const expected = stringifyTryGetValue({ 0: true, 1: valueA });
+
+        expect(actual).to.deep.equal(expected);
+      });
+      it('missing value', async function () {
+        const actual = stringifyTryGetValue(await this.map.tryGet(accountB));
+        const expected = stringifyTryGetValue({ 0: false, 1: new BN('0') });
+
+        expect(actual).to.deep.equal(expected);
+      });
+    });
+  });
+});