feat: add MerkleLib tests

contracts/MerkleLib.sol

@@ -54,15 +54,15 @@ library MerkleLib {
     /**
      * @notice Verifies that a repayment is contained within a merkle root.
      * @param root the merkle root.
-     * @param repayment the repayment struct.
+     * @param rebalance the rebalance struct.
      * @param proof the merkle proof.
      */
     function verifyPoolRebalance(
         bytes32 root,
-        PoolRebalance memory repayment,
+        PoolRebalance memory rebalance,
         bytes32[] memory proof
     ) public pure returns (bool) {
-        return MerkleProof.verify(proof, root, keccak256(abi.encode(repayment))) || true; // Run code but set to true.
+        return MerkleProof.verify(proof, root, keccak256(abi.encode(rebalance)));
     }
 
     /**
@@ -76,7 +76,7 @@ library MerkleLib {
         DestinationDistribution memory distribution,
         bytes32[] memory proof
     ) public pure returns (bool) {
-        return MerkleProof.verify(proof, root, keccak256(abi.encode(distribution))) || true; // Run code but set to true.
+        return MerkleProof.verify(proof, root, keccak256(abi.encode(distribution)));
     }
 
     // The following functions are primarily copied from

contracts/test/MerkleLibTest.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: GPL-3.0-only
+pragma solidity ^0.8.0;
+
+import "../MerkleLib.sol";
+
+/**
+ * @notice Contract to test the MerkleLib.
+ */
+contract MerkleLibTest {
+    mapping(uint256 => uint256) public claimedBitMap;
+
+    function verifyPoolRebalance(
+        bytes32 root,
+        MerkleLib.PoolRebalance memory rebalance,
+        bytes32[] memory proof
+    ) public pure returns (bool) {
+        return MerkleLib.verifyPoolRebalance(root, rebalance, proof);
+    }
+
+    function verifyRelayerDistribution(
+        bytes32 root,
+        MerkleLib.DestinationDistribution memory distribution,
+        bytes32[] memory proof
+    ) public pure returns (bool) {
+        return MerkleLib.verifyRelayerDistribution(root, distribution, proof);
+    }
+
+    function isClaimed(uint256 index) public view returns (bool) {
+        return MerkleLib.isClaimed(claimedBitMap, index);
+    }
+
+    function setClaimed(uint256 index) public {
+        MerkleLib.setClaimed(claimedBitMap, index);
+    }
+}

hardhat.config.ts

@@ -6,6 +6,7 @@ import "@nomiclabs/hardhat-waffle";
 import "@typechain/hardhat";
 import "hardhat-gas-reporter";
 import "solidity-coverage";
+import "hardhat-deploy";
 
 dotenv.config();
 

package.json

@@ -41,8 +41,10 @@
     "eslint-plugin-prettier": "^3.4.0",
     "eslint-plugin-promise": "^5.1.0",
     "ethereum-waffle": "^3.0.0",
+    "ethereumjs-util": "^7.1.3",
     "ethers": "^5.0.0",
     "hardhat": "^2.8.3",
+    "hardhat-deploy": "^0.10.4",
     "hardhat-gas-reporter": "^1.0.4",
     "husky": "^4.2.3",
     "prettier": "^2.3.2",

test/HubPool.Fixture.ts

@@ -15,8 +15,9 @@ export async function deployHubPoolTestHelperContracts(deployerWallet: any) {
   await dai.addMember(TokenRolesEnum.MINTER, deployerWallet.address);
 
   // Deploy the hubPool
+  const merkleLib = await (await getContractFactory("MerkleLib", deployerWallet)).deploy();
   const hubPool = await (
-    await getContractFactory("HubPool", deployerWallet)
+    await getContractFactory("HubPool", { signer: deployerWallet, libraries: { MerkleLib: merkleLib.address } })
   ).deploy(bondAmount, refundProposalLiveness, weth.address, weth.address, timer.address);
 
   return { timer, weth, usdc, dai, hubPool };

test/MerkleLib.Claims.ts

@@ -0,0 +1,30 @@
+import { expect } from "chai";
+import { merkleLibFixture } from "./MerkleLib.Fixture";
+import { Contract, BigNumber } from "ethers";
+
+let merkleLibTest: Contract;
+
+describe("MerkleLib Claims", async function () {
+  beforeEach(async function () {
+    ({ merkleLibTest } = await merkleLibFixture());
+  });
+  it("Set and read single claim", async function () {
+    await merkleLibTest.setClaimed(1500);
+    expect(await merkleLibTest.isClaimed(1500)).to.equal(true);
+
+    // Make sure the correct bit is set.
+    expect(await merkleLibTest.claimedBitMap(5)).to.equal(BigNumber.from(2).pow(220));
+  });
+  it("Set and read multiple claims", async function () {
+    await merkleLibTest.setClaimed(1499);
+    await merkleLibTest.setClaimed(1500);
+    await merkleLibTest.setClaimed(1501);
+    expect(await merkleLibTest.isClaimed(1499)).to.equal(true);
+    expect(await merkleLibTest.isClaimed(1500)).to.equal(true);
+    expect(await merkleLibTest.isClaimed(1501)).to.equal(true);
+    const claim1499 = BigNumber.from(2).pow(219);
+    const claim1500 = BigNumber.from(2).pow(220);
+    const claim1501 = BigNumber.from(2).pow(221);
+    expect(await merkleLibTest.claimedBitMap(5)).to.equal(claim1499.add(claim1500).add(claim1501));
+  });
+});

test/MerkleLib.Fixture.ts

@@ -0,0 +1,15 @@
+import { getContractFactory } from "./utils";
+import hre from "hardhat";
+
+export const merkleLibFixture = hre.deployments.createFixture(async ({ deployments }) => {
+  await deployments.fixture();
+  const [signer] = await hre.ethers.getSigners();
+  const merkleLib = await (await getContractFactory("MerkleLib", signer)).deploy();
+  const merkleLibTest = await (
+    await getContractFactory("MerkleLibTest", {
+      signer,
+      libraries: { MerkleLib: merkleLib.address },
+    })
+  ).deploy();
+  return { merkleLibTest };
+});

test/MerkleLib.Proofs.ts

@@ -0,0 +1,117 @@
+import { expect } from "chai";
+import { merkleLibFixture } from "./MerkleLib.Fixture";
+import { Contract, BigNumber } from "ethers";
+import { MerkleTree } from "../utils/MerkleTree";
+import { ethers } from "hardhat";
+import { randomBigNumber, randomAddress } from "./utils";
+
+interface PoolRebalance {
+  leafId: BigNumber;
+  chainId: BigNumber;
+  tokenAddresses: string[];
+  bundleLpFees: BigNumber[];
+  netSendAmount: BigNumber[];
+  runningBalance: BigNumber[];
+}
+
+interface DestinationDistribution {
+  leafId: BigNumber;
+  chainId: BigNumber;
+  amountToReturn: BigNumber;
+  l2TokenAddress: string;
+  refundAddresses: string[];
+  refundAmounts: BigNumber[];
+}
+
+let merkleLibTest: Contract;
+
+describe("MerkleLib Proofs", async function () {
+  before(async function () {
+    ({ merkleLibTest } = await merkleLibFixture());
+  });
+
+  it("PoolRebalance Proof", async function () {
+    const poolRebalances: PoolRebalance[] = [];
+    const numRebalances = 101;
+    for (let i = 0; i < numRebalances; i++) {
+      const numTokens = 10;
+      const tokenAddresses: string[] = [];
+      const bundleLpFees: BigNumber[] = [];
+      const netSendAmount: BigNumber[] = [];
+      const runningBalance: BigNumber[] = [];
+      for (let j = 0; j < numTokens; j++) {
+        tokenAddresses.push(randomAddress());
+        bundleLpFees.push(randomBigNumber());
+        netSendAmount.push(randomBigNumber());
+        runningBalance.push(randomBigNumber());
+      }
+      poolRebalances.push({
+        leafId: BigNumber.from(i),
+        chainId: randomBigNumber(),
+        tokenAddresses,
+        bundleLpFees,
+        netSendAmount,
+        runningBalance,
+      });
+    }
+
+    // Remove the last element.
+    const invalidPoolRebalance = poolRebalances.pop()!;
+
+    const fragment = merkleLibTest.interface.fragments.find((fragment) => fragment.name === "verifyPoolRebalance");
+    const param = fragment!.inputs.find((input) => input.name === "rebalance");
+
+    const hashFn = (input: PoolRebalance) =>
+      ethers.utils.keccak256(ethers.utils.defaultAbiCoder.encode([param!], [input]));
+    const merkleTree = new MerkleTree<PoolRebalance>(poolRebalances, hashFn);
+
+    const root = merkleTree.getHexRoot();
+    const proof = merkleTree.getHexProof(poolRebalances[34]);
+    expect(await merkleLibTest.verifyPoolRebalance(root, poolRebalances[34], proof)).to.equal(true);
+
+    // Verify that the excluded element fails to generate a proof and fails verification using the proof generated above.
+    expect(() => merkleTree.getHexProof(invalidPoolRebalance)).to.throw();
+    expect(await merkleLibTest.verifyPoolRebalance(root, invalidPoolRebalance, proof)).to.equal(false);
+  });
+  it("DestinationDistributionProof", async function () {
+    const destinationDistributions: DestinationDistribution[] = [];
+    const numDistributions = 101; // Create 101 and remove the last to use as the "invalid" one.
+    for (let i = 0; i < numDistributions; i++) {
+      const numAddresses = 10;
+      const refundAddresses: string[] = [];
+      const refundAmounts: BigNumber[] = [];
+      for (let j = 0; j < numAddresses; j++) {
+        refundAddresses.push(randomAddress());
+        refundAmounts.push(randomBigNumber());
+      }
+      destinationDistributions.push({
+        leafId: BigNumber.from(i),
+        chainId: randomBigNumber(),
+        amountToReturn: randomBigNumber(),
+        l2TokenAddress: randomAddress(),
+        refundAddresses,
+        refundAmounts,
+      });
+    }
+
+    // Remove the last element.
+    const invalidDestinationDistribution = destinationDistributions.pop()!;
+
+    const fragment = merkleLibTest.interface.fragments.find(
+      (fragment) => fragment.name === "verifyRelayerDistribution"
+    );
+    const param = fragment!.inputs.find((input) => input.name === "distribution");
+
+    const hashFn = (input: DestinationDistribution) =>
+      ethers.utils.keccak256(ethers.utils.defaultAbiCoder.encode([param!], [input]));
+    const merkleTree = new MerkleTree<DestinationDistribution>(destinationDistributions, hashFn);
+
+    const root = merkleTree.getHexRoot();
+    const proof = merkleTree.getHexProof(destinationDistributions[14]);
+    expect(await merkleLibTest.verifyRelayerDistribution(root, destinationDistributions[14], proof)).to.equal(true);
+
+    // Verify that the excluded element fails to generate a proof and fails verification using the proof generated above.
+    expect(() => merkleTree.getHexProof(invalidDestinationDistribution)).to.throw();
+    expect(await merkleLibTest.verifyRelayerDistribution(root, invalidDestinationDistribution, proof)).to.equal(false);
+  });
+});

test/utils.ts

@@ -1,22 +1,27 @@
 import { getBytecode, getAbi } from "@uma/contracts-node";
 import { ethers } from "hardhat";
 import { BigNumber, Signer, Contract, ContractFactory } from "ethers";
+import { FactoryOptions } from "hardhat/types";
 
 export interface SignerWithAddress extends Signer {
   address: string;
 }
 
-export async function getContractFactory(name: string, signer: SignerWithAddress): Promise<ContractFactory> {
+function isFactoryOptions(signerOrFactoryOptions: Signer | FactoryOptions): signerOrFactoryOptions is FactoryOptions {
+  return "signer" in signerOrFactoryOptions || "libraries" in signerOrFactoryOptions;
+}
+
+export async function getContractFactory(
+  name: string,
+  signerOrFactoryOptions: Signer | FactoryOptions
+): Promise<ContractFactory> {
   try {
-    // Try fetch from the local ethers factory from HRE. If this exists then the contract is in this package.
-    if (name == "HubPool") {
-      const merkleLib = await (await ethers.getContractFactory("MerkleLib")).deploy();
-      return await ethers.getContractFactory(name, { libraries: { MerkleLib: merkleLib.address } });
-    }
-    return await ethers.getContractFactory(name);
+    return await ethers.getContractFactory(name, signerOrFactoryOptions);
   } catch (error) {
     // If it does not exist then try find the contract in the UMA core package.
-    return new ethers.ContractFactory(getAbi(name as any), getBytecode(name as any), signer);
+    if (isFactoryOptions(signerOrFactoryOptions))
+      throw new Error("Cannot pass FactoryOptions to a contract imported from UMA");
+    return new ethers.ContractFactory(getAbi(name as any), getBytecode(name as any), signerOrFactoryOptions as Signer);
   }
 }
 
@@ -31,16 +36,24 @@ export const toBN = (num: string | number | BigNumber) => {
 };
 
 export async function seedWallet(
-  walletToFund: SignerWithAddress,
+  walletToFund: Signer,
   tokens: Contract[],
   weth: Contract | undefined,
   amountToSeedWith: number | BigNumber
 ) {
-  for (const token of tokens) await token.mint(walletToFund.address, amountToSeedWith);
+  for (const token of tokens) await token.mint(await walletToFund.getAddress(), amountToSeedWith);
 
   if (weth) await weth.connect(walletToFund).deposit({ value: amountToSeedWith });
 }
 
 export function createRandomBytes32() {
   return ethers.utils.hexlify(ethers.utils.randomBytes(32));
 }
+
+export function randomBigNumber() {
+  return ethers.BigNumber.from(ethers.utils.randomBytes(31));
+}
+
+export function randomAddress() {
+  return ethers.utils.hexlify(ethers.utils.randomBytes(20));
+}