From 1b938e39a82bfea896f45408113f9ef74284461f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=A1s=20Venturo?= Date: Tue, 4 Feb 2020 19:15:32 -0300 Subject: [PATCH] EnumerableSet improvements (#2077) * Remove newAddressSet * Add count and get functions. * Fix lint (cherry picked from commit 7988c044e075024e6018527fddd83c70cc6ec8f0) --- contracts/mocks/EnumerableSetMock.sol | 12 +++++--- contracts/utils/EnumerableSet.sol | 43 ++++++++++++++++++++------- test/utils/EnumerableSet.test.js | 36 +++++++++++++--------- 3 files changed, 62 insertions(+), 29 deletions(-) diff --git a/contracts/mocks/EnumerableSetMock.sol b/contracts/mocks/EnumerableSetMock.sol index 4ef4fa04a15..d2dc13d1ef1 100644 --- a/contracts/mocks/EnumerableSetMock.sol +++ b/contracts/mocks/EnumerableSetMock.sol @@ -9,10 +9,6 @@ contract EnumerableSetMock{ EnumerableSet.AddressSet private set; - constructor() public { - set = EnumerableSet.newAddressSet(); - } - function contains(address value) public view returns (bool) { return set.contains(value); } @@ -30,4 +26,12 @@ contract EnumerableSetMock{ function enumerate() public view returns (address[] memory) { return set.enumerate(); } + + function length() public view returns (uint256) { + return set.length(); + } + + function get(uint256 index) public view returns (address) { + return set.get(index); + } } diff --git a/contracts/utils/EnumerableSet.sol b/contracts/utils/EnumerableSet.sol index 6798a30df84..8ed12155c23 100644 --- a/contracts/utils/EnumerableSet.sol +++ b/contracts/utils/EnumerableSet.sol @@ -29,17 +29,6 @@ library EnumerableSet { address[] values; } - /** - * @dev Creates a new empty address set. - */ - function newAddressSet() - internal - pure - returns (AddressSet memory) - { - return AddressSet({values: new address[](0)}); - } - /** * @dev Add a value to a set. O(1). * Returns false if the value was already in the set. @@ -105,6 +94,9 @@ library EnumerableSet { * @dev Returns an array with all values in the set. O(N). * 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. + + * WARNING: This function may run out of gas on large sets: use {length} and + * {get} instead in these cases. */ function enumerate(AddressSet storage set) internal @@ -117,4 +109,33 @@ library EnumerableSet { } return output; } + + /** + * @dev Returns the number of elements on 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. + */ + function length(AddressSet storage set) + internal + view + returns (uint256) + { + return set.values.length; + } + + /** @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 get(AddressSet storage set, uint256 index) + internal + view + returns (address) + { + return set.values[index]; + } } diff --git a/test/utils/EnumerableSet.test.js b/test/utils/EnumerableSet.test.js index 0162da5fa85..137c8051820 100644 --- a/test/utils/EnumerableSet.test.js +++ b/test/utils/EnumerableSet.test.js @@ -11,29 +11,39 @@ describe('EnumerableSet', function () { this.set = await EnumerableSetMock.new(); }); + async function expectMembersMatch (set, members) { + await Promise.all(members.map(async account => + expect(await set.contains(account)).to.equal(true) + )); + + expect(await set.enumerate()).to.have.same.members(members); + + expect(await set.length()).to.bignumber.equal(members.length.toString()); + + expect(await Promise.all([...Array(members.length).keys()].map(index => + set.get(index) + ))).to.have.same.members(members); + } + it('starts empty', async function () { expect(await this.set.contains(accountA)).to.equal(false); - expect(await this.set.enumerate()).to.have.same.members([]); + + await expectMembersMatch(this.set, []); }); it('adds a value', async function () { const receipt = await this.set.add(accountA); expectEvent(receipt, 'TransactionResult', { result: true }); - expect(await this.set.contains(accountA)).to.equal(true); - expect(await this.set.enumerate()).to.have.same.members([ accountA ]); + await expectMembersMatch(this.set, [accountA]); }); it('adds several values', async function () { await this.set.add(accountA); await this.set.add(accountB); - expect(await this.set.contains(accountA)).to.equal(true); - expect(await this.set.contains(accountB)).to.equal(true); - + await expectMembersMatch(this.set, [accountA, accountB]); expect(await this.set.contains(accountC)).to.equal(false); - - expect(await this.set.enumerate()).to.have.same.members([ accountA, accountB ]); }); it('returns false when adding elements already in the set', async function () { @@ -42,7 +52,7 @@ describe('EnumerableSet', function () { const receipt = (await this.set.add(accountA)); expectEvent(receipt, 'TransactionResult', { result: false }); - expect(await this.set.enumerate()).to.have.same.members([ accountA ]); + await expectMembersMatch(this.set, [accountA]); }); it('removes added values', async function () { @@ -52,7 +62,7 @@ describe('EnumerableSet', function () { expectEvent(receipt, 'TransactionResult', { result: true }); expect(await this.set.contains(accountA)).to.equal(false); - expect(await this.set.enumerate()).to.have.same.members([]); + await expectMembersMatch(this.set, []); }); it('returns false when removing elements not in the set', async function () { @@ -99,10 +109,8 @@ describe('EnumerableSet', function () { // [A, C] - expect(await this.set.contains(accountA)).to.equal(true); - expect(await this.set.contains(accountB)).to.equal(false); - expect(await this.set.contains(accountC)).to.equal(true); + await expectMembersMatch(this.set, [accountA, accountC]); - expect(await this.set.enumerate()).to.have.same.members([ accountA, accountC ]); + expect(await this.set.contains(accountB)).to.equal(false); }); });