Add AddressToUintMap (#3150)

* add AddressToUintMap

* Update contracts/utils/structs/EnumerableMap.sol

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>

* address comments

* lint code

* merge mocks into a single file

* add PR link to changelog entry

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>
Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>

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

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

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