Refactor enumerableMap generate and tests (#4760)

scripts/generate/templates/EnumerableMap.js

@@ -1,12 +1,6 @@
 const format = require('../format-lines');
 const { fromBytes32, toBytes32 } = require('./conversion');
-
-const TYPES = [
-  { name: 'UintToUintMap', keyType: 'uint256', valueType: 'uint256' },
-  { name: 'UintToAddressMap', keyType: 'uint256', valueType: 'address' },
-  { name: 'AddressToUintMap', keyType: 'address', valueType: 'uint256' },
-  { name: 'Bytes32ToUintMap', keyType: 'bytes32', valueType: 'uint256' },
-];
+const { TYPES } = require('./EnumerableMap.opts');
 
 /* eslint-disable max-len */
 const header = `\

scripts/generate/templates/EnumerableMap.opts.js

@@ -0,0 +1,18 @@
+const mapType = str => (str == 'uint256' ? 'Uint' : `${str.charAt(0).toUpperCase()}${str.slice(1)}`);
+const formatType = (keyType, valueType) => ({
+  name: `${mapType(keyType)}To${mapType(valueType)}Map`,
+  keyType,
+  valueType,
+});
+
+const TYPES = [
+  ['uint256', 'uint256'],
+  ['uint256', 'address'],
+  ['address', 'uint256'],
+  ['bytes32', 'uint256'],
+].map(args => formatType(...args));
+
+module.exports = {
+  TYPES,
+  formatType,
+};

test/utils/structs/EnumerableMap.behavior.js

@@ -3,7 +3,7 @@ const { ethers } = require('hardhat');
 
 const zip = (array1, array2) => array1.map((item, index) => [item, array2[index]]);
 
-function shouldBehaveLikeMap(zeroValue, keyType, events) {
+function shouldBehaveLikeMap() {
   async function expectMembersMatch(methods, keys, values) {
     expect(keys.length).to.equal(values.length);
     expect(await methods.length()).to.equal(keys.length);
@@ -25,7 +25,7 @@ function shouldBehaveLikeMap(zeroValue, keyType, events) {
 
   describe('set', function () {
     it('adds a key', async function () {
-      await expect(this.methods.set(this.keyA, this.valueA)).to.emit(this.mock, events.setReturn).withArgs(true);
+      await expect(this.methods.set(this.keyA, this.valueA)).to.emit(this.mock, this.events.setReturn).withArgs(true);
 
       await expectMembersMatch(this.methods, [this.keyA], [this.valueA]);
     });
@@ -41,7 +41,7 @@ function shouldBehaveLikeMap(zeroValue, keyType, events) {
     it('returns false when adding keys already in the set', async function () {
       await this.methods.set(this.keyA, this.valueA);
 
-      await expect(this.methods.set(this.keyA, this.valueA)).to.emit(this.mock, events.setReturn).withArgs(false);
+      await expect(this.methods.set(this.keyA, this.valueA)).to.emit(this.mock, this.events.setReturn).withArgs(false);
 
       await expectMembersMatch(this.methods, [this.keyA], [this.valueA]);
     });
@@ -58,7 +58,7 @@ function shouldBehaveLikeMap(zeroValue, keyType, events) {
     it('removes added keys', async function () {
       await this.methods.set(this.keyA, this.valueA);
 
-      await expect(this.methods.remove(this.keyA)).to.emit(this.mock, events.removeReturn).withArgs(true);
+      await expect(this.methods.remove(this.keyA)).to.emit(this.mock, this.events.removeReturn).withArgs(true);
 
       expect(await this.methods.contains(this.keyA)).to.be.false;
       await expectMembersMatch(this.methods, [], []);
@@ -66,7 +66,7 @@ function shouldBehaveLikeMap(zeroValue, keyType, events) {
 
     it('returns false when removing keys not in the set', async function () {
       await expect(await this.methods.remove(this.keyA))
-        .to.emit(this.mock, events.removeReturn)
+        .to.emit(this.mock, this.events.removeReturn)
         .withArgs(false);
 
       expect(await this.methods.contains(this.keyA)).to.be.false;
@@ -130,7 +130,7 @@ function shouldBehaveLikeMap(zeroValue, keyType, events) {
       it('missing value', async function () {
         await expect(this.methods.get(this.keyB))
           .to.be.revertedWithCustomError(this.mock, 'EnumerableMapNonexistentKey')
-          .withArgs(ethers.AbiCoder.defaultAbiCoder().encode([keyType], [this.keyB]));
+          .withArgs(ethers.AbiCoder.defaultAbiCoder().encode([this.keyType], [this.keyB]));
       });
     });
 
@@ -140,7 +140,7 @@ function shouldBehaveLikeMap(zeroValue, keyType, events) {
       });
 
       it('missing value', async function () {
-        expect(await this.methods.tryGet(this.keyB)).to.have.ordered.members([false, zeroValue]);
+        expect(await this.methods.tryGet(this.keyB)).to.have.ordered.members([false, this.zeroValue]);
       });
     });
   });

test/utils/structs/EnumerableMap.test.js

@@ -2,6 +2,7 @@ const { ethers } = require('hardhat');
 const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers');
 const { mapValues } = require('../../helpers/iterate');
 const { randomArray, generators } = require('../../helpers/random');
+const { TYPES, formatType } = require('../../../scripts/generate/templates/EnumerableMap.opts');
 
 const { shouldBehaveLikeMap } = require('./EnumerableMap.behavior');
 
@@ -14,132 +15,77 @@ const getMethods = (mock, fnSigs) => {
   );
 };
 
-describe('EnumerableMap', function () {
-  // UintToAddressMap
-  describe('UintToAddressMap', function () {
-    const fixture = async () => {
-      const mock = await ethers.deployContract('$EnumerableMap');
-
-      const [keyA, keyB, keyC] = randomArray(generators.uint256);
-      const [valueA, valueB, valueC] = randomArray(generators.address);
-
-      const methods = getMethods(mock, {
-        set: '$set(uint256,uint256,address)',
-        get: '$get_EnumerableMap_UintToAddressMap(uint256,uint256)',
-        tryGet: '$tryGet_EnumerableMap_UintToAddressMap(uint256,uint256)',
-        remove: '$remove_EnumerableMap_UintToAddressMap(uint256,uint256)',
-        length: '$length_EnumerableMap_UintToAddressMap(uint256)',
-        at: '$at_EnumerableMap_UintToAddressMap(uint256,uint256)',
-        contains: '$contains_EnumerableMap_UintToAddressMap(uint256,uint256)',
-        keys: '$keys_EnumerableMap_UintToAddressMap(uint256)',
-      });
-
-      return { mock, keyA, keyB, keyC, valueA, valueB, valueC, methods };
-    };
-
-    beforeEach(async function () {
-      Object.assign(this, await loadFixture(fixture));
-    });
-
-    shouldBehaveLikeMap(ethers.ZeroAddress, 'uint256', {
-      setReturn: 'return$set_EnumerableMap_UintToAddressMap_uint256_address',
-      removeReturn: 'return$remove_EnumerableMap_UintToAddressMap_uint256',
-    });
-  });
-
-  // Bytes32ToBytes32Map
-  describe('Bytes32ToBytes32Map', function () {
-    const fixture = async () => {
-      const mock = await ethers.deployContract('$EnumerableMap');
-
-      const [keyA, keyB, keyC] = randomArray(generators.bytes32);
-      const [valueA, valueB, valueC] = randomArray(generators.bytes32);
-
-      const methods = getMethods(mock, {
-        set: '$set(uint256,bytes32,bytes32)',
-        get: '$get_EnumerableMap_Bytes32ToBytes32Map(uint256,bytes32)',
-        tryGet: '$tryGet_EnumerableMap_Bytes32ToBytes32Map(uint256,bytes32)',
-        remove: '$remove_EnumerableMap_Bytes32ToBytes32Map(uint256,bytes32)',
-        length: '$length_EnumerableMap_Bytes32ToBytes32Map(uint256)',
-        at: '$at_EnumerableMap_Bytes32ToBytes32Map(uint256,uint256)',
-        contains: '$contains_EnumerableMap_Bytes32ToBytes32Map(uint256,bytes32)',
-        keys: '$keys_EnumerableMap_Bytes32ToBytes32Map(uint256)',
-      });
-
-      return { mock, keyA, keyB, keyC, valueA, valueB, valueC, methods };
-    };
-
-    beforeEach(async function () {
-      Object.assign(this, await loadFixture(fixture));
-    });
-
-    shouldBehaveLikeMap(ethers.ZeroHash, 'bytes32', {
-      setReturn: 'return$set_EnumerableMap_Bytes32ToBytes32Map_bytes32_bytes32',
-      removeReturn: 'return$remove_EnumerableMap_Bytes32ToBytes32Map_bytes32',
-    });
-  });
-
-  // UintToUintMap
-  describe('UintToUintMap', function () {
-    const fixture = async () => {
-      const mock = await ethers.deployContract('$EnumerableMap');
-
-      const [keyA, keyB, keyC] = randomArray(generators.uint256);
-      const [valueA, valueB, valueC] = randomArray(generators.uint256);
-
-      const methods = getMethods(mock, {
-        set: '$set(uint256,uint256,uint256)',
-        get: '$get_EnumerableMap_UintToUintMap(uint256,uint256)',
-        tryGet: '$tryGet_EnumerableMap_UintToUintMap(uint256,uint256)',
-        remove: '$remove_EnumerableMap_UintToUintMap(uint256,uint256)',
-        length: '$length_EnumerableMap_UintToUintMap(uint256)',
-        at: '$at_EnumerableMap_UintToUintMap(uint256,uint256)',
-        contains: '$contains_EnumerableMap_UintToUintMap(uint256,uint256)',
-        keys: '$keys_EnumerableMap_UintToUintMap(uint256)',
-      });
-
-      return { mock, keyA, keyB, keyC, valueA, valueB, valueC, methods };
-    };
+const testTypes = [formatType('bytes32', 'bytes32'), ...TYPES];
+
+async function fixture() {
+  const mock = await ethers.deployContract('$EnumerableMap');
+
+  const zeroValue = {
+    uint256: 0n,
+    address: ethers.ZeroAddress,
+    bytes32: ethers.ZeroHash,
+  };
+
+  const env = Object.fromEntries(
+    testTypes.map(({ name, keyType, valueType }) => [
+      name,
+      {
+        keyType,
+        keys: randomArray(generators[keyType]),
+        values: randomArray(generators[valueType]),
+
+        methods: getMethods(
+          mock,
+          testTypes.filter(t => keyType == t.keyType).length == 1
+            ? {
+                set: `$set(uint256,${keyType},${valueType})`,
+                get: `$get(uint256,${keyType})`,
+                tryGet: `$tryGet(uint256,${keyType})`,
+                remove: `$remove(uint256,${keyType})`,
+                length: `$length_EnumerableMap_${name}(uint256)`,
+                at: `$at_EnumerableMap_${name}(uint256,uint256)`,
+                contains: `$contains(uint256,${keyType})`,
+                keys: `$keys_EnumerableMap_${name}(uint256)`,
+              }
+            : {
+                set: `$set(uint256,${keyType},${valueType})`,
+                get: `$get_EnumerableMap_${name}(uint256,${keyType})`,
+                tryGet: `$tryGet_EnumerableMap_${name}(uint256,${keyType})`,
+                remove: `$remove_EnumerableMap_${name}(uint256,${keyType})`,
+                length: `$length_EnumerableMap_${name}(uint256)`,
+                at: `$at_EnumerableMap_${name}(uint256,uint256)`,
+                contains: `$contains_EnumerableMap_${name}(uint256,${keyType})`,
+                keys: `$keys_EnumerableMap_${name}(uint256)`,
+              },
+        ),
+
+        zeroValue: zeroValue[valueType],
+        events: {
+          setReturn: `return$set_EnumerableMap_${name}_${keyType}_${valueType}`,
+          removeReturn: `return$remove_EnumerableMap_${name}_${keyType}`,
+        },
+      },
+    ]),
+  );
 
-    beforeEach(async function () {
-      Object.assign(this, await loadFixture(fixture));
-    });
+  return { mock, env };
+}
 
-    shouldBehaveLikeMap(0n, 'uint256', {
-      setReturn: 'return$set_EnumerableMap_UintToUintMap_uint256_uint256',
-      removeReturn: 'return$remove_EnumerableMap_UintToUintMap_uint256',
-    });
+describe('EnumerableMap', function () {
+  beforeEach(async function () {
+    Object.assign(this, await loadFixture(fixture));
   });
 
-  // Bytes32ToUintMap
-  describe('Bytes32ToUintMap', function () {
-    const fixture = async () => {
-      const mock = await ethers.deployContract('$EnumerableMap');
-
-      const [keyA, keyB, keyC] = randomArray(generators.bytes32);
-      const [valueA, valueB, valueC] = randomArray(generators.uint256);
-
-      const methods = getMethods(mock, {
-        set: '$set(uint256,bytes32,uint256)',
-        get: '$get_EnumerableMap_Bytes32ToUintMap(uint256,bytes32)',
-        tryGet: '$tryGet_EnumerableMap_Bytes32ToUintMap(uint256,bytes32)',
-        remove: '$remove_EnumerableMap_Bytes32ToUintMap(uint256,bytes32)',
-        length: '$length_EnumerableMap_Bytes32ToUintMap(uint256)',
-        at: '$at_EnumerableMap_Bytes32ToUintMap(uint256,uint256)',
-        contains: '$contains_EnumerableMap_Bytes32ToUintMap(uint256,bytes32)',
-        keys: '$keys_EnumerableMap_Bytes32ToUintMap(uint256)',
+  // UintToAddressMap
+  for (const { name } of testTypes) {
+    describe(name, function () {
+      beforeEach(async function () {
+        Object.assign(this, this.env[name]);
+        [this.keyA, this.keyB, this.keyC] = this.keys;
+        [this.valueA, this.valueB, this.valueC] = this.values;
       });
 
-      return { mock, keyA, keyB, keyC, valueA, valueB, valueC, methods };
-    };
-
-    beforeEach(async function () {
-      Object.assign(this, await loadFixture(fixture));
-    });
-
-    shouldBehaveLikeMap(0n, 'bytes32', {
-      setReturn: 'return$set_EnumerableMap_Bytes32ToUintMap_bytes32_uint256',
-      removeReturn: 'return$remove_EnumerableMap_Bytes32ToUintMap_bytes32',
+      shouldBehaveLikeMap();
     });
-  });
+  }
 });