EnumerableSet improvements (#2077)

* Remove newAddressSet

* Add count and get functions.

* Fix lint

(cherry picked from commit 7988c044)

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);
+    }
 }

contracts/utils/EnumerableSet.sol

@@ -30,17 +30,6 @@ library EnumerableSet {
     }
 
     /**
-     * @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];
+    }
 }

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);
   });
 });