diff --git a/.prettierrc b/.prettierrc index 5c11cc2237d..f91ad7ee614 100644 --- a/.prettierrc +++ b/.prettierrc @@ -1,8 +1,11 @@ { + "singleQuote": true, + "trailingComma": "all", "overrides": [ { "files": "*.sol", "options": { + "singleQuote": false, "printWidth": 120, "explicitTypes": "always" } diff --git a/CHANGELOG.md b/CHANGELOG.md index e878c2c43b3..0bb0981f52f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## Unreleased * `AccessControl`: add a virtual `_checkRole(bytes32)` function that can be overriden to alter the `onlyRole` modifier behavior. ([#3137](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3137)) + * `EnumerableMap`: add new `AddressToUintMap` map type. ([#3150](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3150)) ## Unreleased diff --git a/contracts/mocks/EnumerableMapMock.sol b/contracts/mocks/EnumerableMapMock.sol index 510647b5839..21c1fe362d9 100644 --- a/contracts/mocks/EnumerableMapMock.sol +++ b/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); + } +} diff --git a/contracts/utils/structs/EnumerableMap.sol b/contracts/utils/structs/EnumerableMap.sol index 1bc571594cb..1c3f4357ecf 100644 --- a/contracts/utils/structs/EnumerableMap.sol +++ b/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))))); + } } diff --git a/test/utils/structs/EnumerableMap/AddressToUintMap.test.js b/test/utils/structs/EnumerableMap/AddressToUintMap.test.js new file mode 100644 index 00000000000..0a79479323e --- /dev/null +++ b/test/utils/structs/EnumerableMap/AddressToUintMap.test.js @@ -0,0 +1,157 @@ +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('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); + }); + }); + }); +}); diff --git a/test/utils/structs/EnumerableMap.test.js b/test/utils/structs/EnumerableMap/UintToAddressMap.test.js similarity index 78% rename from test/utils/structs/EnumerableMap.test.js rename to test/utils/structs/EnumerableMap/UintToAddressMap.test.js index 9dc700b510b..560ed507701 100644 --- a/test/utils/structs/EnumerableMap.test.js +++ b/test/utils/structs/EnumerableMap/UintToAddressMap.test.js @@ -1,44 +1,20 @@ const { BN, constants, expectEvent, expectRevert } = require('@openzeppelin/test-helpers'); const { expect } = require('chai'); +const { expectMembersMatch } = require('./helpers'); -const zip = require('lodash.zip'); +const UintToAddressMapMock = artifacts.require('UintToAddressMapMock'); -const EnumerableMapMock = artifacts.require('EnumerableMapMock'); - -contract('EnumerableMap', function (accounts) { - const [ accountA, accountB, accountC ] = accounts; +contract('UintToAddressMap', function (accounts) { + const [accountA, accountB, accountC] = accounts; const keyA = new BN('7891'); const keyB = new BN('451'); const keyC = new BN('9592328'); beforeEach(async function () { - this.map = await EnumerableMapMock.new(); + this.map = await UintToAddressMapMock.new(); }); - async function expectMembersMatch (map, keys, values) { - expect(keys.length).to.equal(values.length); - - await Promise.all(keys.map(async key => - expect(await map.contains(key)).to.equal(true), - )); - - expect(await map.length()).to.bignumber.equal(keys.length.toString()); - - expect(await Promise.all(keys.map(key => - map.get(key), - ))).to.have.same.members(values); - - // To compare key-value pairs, we zip keys and values, and convert BNs to - // strings to workaround Chai limitations when dealing with nested arrays - expect(await Promise.all([...Array(keys.length).keys()].map(async (index) => { - const entry = await map.at(index); - return [entry.key.toString(), entry.value]; - }))).to.have.same.deep.members( - zip(keys.map(k => k.toString()), values), - ); - } - it('starts empty', async function () { expect(await this.map.contains(keyA)).to.equal(false); @@ -64,7 +40,7 @@ contract('EnumerableMap', function (accounts) { it('returns false when adding keys already in the set', async function () { await this.map.set(keyA, accountA); - const receipt = (await this.map.set(keyA, accountA)); + const receipt = await this.map.set(keyA, accountA); expectEvent(receipt, 'OperationResult', { result: false }); await expectMembersMatch(this.map, [keyA], [accountA]); diff --git a/test/utils/structs/EnumerableMap/helpers.js b/test/utils/structs/EnumerableMap/helpers.js new file mode 100644 index 00000000000..456cdf156f3 --- /dev/null +++ b/test/utils/structs/EnumerableMap/helpers.js @@ -0,0 +1,31 @@ +const zip = require('lodash.zip'); + +const toStringArray = (array) => array.map((i) => i.toString()); + +async function expectMembersMatch (map, keys, values) { + const stringKeys = toStringArray(keys); + const stringValues = toStringArray(values); + + expect(keys.length).to.equal(values.length); + + await Promise.all(keys.map(async (key) => expect(await map.contains(key)).to.equal(true))); + + expect(await map.length()).to.bignumber.equal(keys.length.toString()); + + expect(toStringArray(await Promise.all(keys.map((key) => map.get(key))))).to.have.same.members(stringValues); + + // to compare key-value pairs, we zip keys and values + // we also stringify pairs because this helper is used for multiple types of maps + expect( + await Promise.all( + [...Array(keys.length).keys()].map(async (index) => { + const { key, value } = await map.at(index); + return [key.toString(), value.toString()]; + }), + ), + ).to.have.same.deep.members(zip(stringKeys, stringValues)); +} + +module.exports = { + expectMembersMatch, +};