Add EnumerableMap, refactor ERC721 (#2160)

* Implement AddressSet in terms of a generic Set

* Add Uint256Set

* Add EnumerableMap

* Fix wording on EnumerableSet docs and tests

* Refactor ERC721 using EnumerableSet and EnumerableMap

* Fix tests

* Fix linter error

* Gas optimization for EnumerableMap

* Gas optimization for EnumerableSet

* Remove often not-taken if from Enumerable data structures

* Fix failing test

* Gas optimization for EnumerableMap

* Fix linter errors

* Add comment for clarification

* Improve test naming

* Rename EnumerableMap.add to set

* Add overload for EnumerableMap.get with custom error message

* Improve Enumerable docs

* Rename Uint256Set to UintSet

* Add changelog entry

test/utils/EnumerableMap.test.js

@@ -0,0 +1,139 @@
+const { accounts, contract } = require('@openzeppelin/test-environment');
+const { BN, expectEvent } = require('@openzeppelin/test-helpers');
+const { expect } = require('chai');
+
+const zip = require('lodash.zip');
+
+const EnumerableMapMock = contract.fromArtifact('EnumerableMapMock');
+
+describe('EnumerableMap', function () {
+  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();
+  });
+
+  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);
+
+    await expectMembersMatch(this.map, [], []);
+  });
+
+  it('adds a key', async function () {
+    const receipt = await this.map.set(keyA, accountA);
+    expectEvent(receipt, 'OperationResult', { result: true });
+
+    await expectMembersMatch(this.map, [keyA], [accountA]);
+  });
+
+  it('adds several keys', async function () {
+    await this.map.set(keyA, accountA);
+    await this.map.set(keyB, accountB);
+
+    await expectMembersMatch(this.map, [keyA, keyB], [accountA, accountB]);
+    expect(await this.map.contains(keyC)).to.equal(false);
+  });
+
+  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));
+    expectEvent(receipt, 'OperationResult', { result: false });
+
+    await expectMembersMatch(this.map, [keyA], [accountA]);
+  });
+
+  it('updates values for keys already in the set', async function () {
+    await this.map.set(keyA, accountA);
+
+    await this.map.set(keyA, accountB);
+
+    await expectMembersMatch(this.map, [keyA], [accountB]);
+  });
+
+  it('removes added keys', async function () {
+    await this.map.set(keyA, accountA);
+
+    const receipt = await this.map.remove(keyA);
+    expectEvent(receipt, 'OperationResult', { result: true });
+
+    expect(await this.map.contains(keyA)).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(keyA);
+    expectEvent(receipt, 'OperationResult', { result: false });
+
+    expect(await this.map.contains(keyA)).to.equal(false);
+  });
+
+  it('adds and removes multiple keys', async function () {
+    // []
+
+    await this.map.set(keyA, accountA);
+    await this.map.set(keyC, accountC);
+
+    // [A, C]
+
+    await this.map.remove(keyA);
+    await this.map.remove(keyB);
+
+    // [C]
+
+    await this.map.set(keyB, accountB);
+
+    // [C, B]
+
+    await this.map.set(keyA, accountA);
+    await this.map.remove(keyC);
+
+    // [A, B]
+
+    await this.map.set(keyA, accountA);
+    await this.map.set(keyB, accountB);
+
+    // [A, B]
+
+    await this.map.set(keyC, accountC);
+    await this.map.remove(keyA);
+
+    // [B, C]
+
+    await this.map.set(keyA, accountA);
+    await this.map.remove(keyB);
+
+    // [A, C]
+
+    await expectMembersMatch(this.map, [keyA, keyC], [accountA, accountC]);
+
+    expect(await this.map.contains(keyB)).to.equal(false);
+  });
+});

test/utils/EnumerableSet.test.js

@@ -11,18 +11,16 @@ describe('EnumerableSet', function () {
     this.set = await EnumerableSetMock.new();
   });
 
-  async function expectMembersMatch (set, members) {
-    await Promise.all(members.map(async account =>
+  async function expectMembersMatch (set, values) {
+    await Promise.all(values.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(values.length.toString());
 
-    expect(await set.length()).to.bignumber.equal(members.length.toString());
-
-    expect(await Promise.all([...Array(members.length).keys()].map(index =>
+    expect(await Promise.all([...Array(values.length).keys()].map(index =>
       set.at(index)
-    ))).to.have.same.members(members);
+    ))).to.have.same.members(values);
   }
 
   it('starts empty', async function () {
@@ -33,7 +31,7 @@ describe('EnumerableSet', function () {
 
   it('adds a value', async function () {
     const receipt = await this.set.add(accountA);
-    expectEvent(receipt, 'TransactionResult', { result: true });
+    expectEvent(receipt, 'OperationResult', { result: true });
 
     await expectMembersMatch(this.set, [accountA]);
   });
@@ -46,11 +44,11 @@ describe('EnumerableSet', function () {
     expect(await this.set.contains(accountC)).to.equal(false);
   });
 
-  it('returns false when adding elements already in the set', async function () {
+  it('returns false when adding values already in the set', async function () {
     await this.set.add(accountA);
 
     const receipt = (await this.set.add(accountA));
-    expectEvent(receipt, 'TransactionResult', { result: false });
+    expectEvent(receipt, 'OperationResult', { result: false });
 
     await expectMembersMatch(this.set, [accountA]);
   });
@@ -63,15 +61,15 @@ describe('EnumerableSet', function () {
     await this.set.add(accountA);
 
     const receipt = await this.set.remove(accountA);
-    expectEvent(receipt, 'TransactionResult', { result: true });
+    expectEvent(receipt, 'OperationResult', { result: true });
 
     expect(await this.set.contains(accountA)).to.equal(false);
     await expectMembersMatch(this.set, []);
   });
 
-  it('returns false when removing elements not in the set', async function () {
+  it('returns false when removing values not in the set', async function () {
     const receipt = await this.set.remove(accountA);
-    expectEvent(receipt, 'TransactionResult', { result: false });
+    expectEvent(receipt, 'OperationResult', { result: false });
 
     expect(await this.set.contains(accountA)).to.equal(false);
   });