Add a MerkleTree builder (#3617)

Co-authored-by: Ernesto García <ernestognw@gmail.com>

test/utils/structs/MerkleTree.test.js

@@ -0,0 +1,100 @@
+const { ethers } = require('hardhat');
+const { expect } = require('chai');
+const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers');
+const { PANIC_CODES } = require('@nomicfoundation/hardhat-chai-matchers/panic');
+const { StandardMerkleTree } = require('@openzeppelin/merkle-tree');
+
+const { generators } = require('../../helpers/random');
+
+const makeTree = (leafs = [ethers.ZeroHash]) =>
+  StandardMerkleTree.of(
+    leafs.map(leaf => [leaf]),
+    ['bytes32'],
+    { sortLeaves: false },
+  );
+
+const hashLeaf = leaf => makeTree().leafHash([leaf]);
+
+const DEPTH = 4n; // 16 slots
+const ZERO = hashLeaf(ethers.ZeroHash);
+
+async function fixture() {
+  const mock = await ethers.deployContract('MerkleTreeMock');
+  await mock.setup(DEPTH, ZERO);
+  return { mock };
+}
+
+describe('MerkleTree', function () {
+  beforeEach(async function () {
+    Object.assign(this, await loadFixture(fixture));
+  });
+
+  it('sets initial values at setup', async function () {
+    const merkleTree = makeTree(Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash));
+
+    expect(await this.mock.root()).to.equal(merkleTree.root);
+    expect(await this.mock.depth()).to.equal(DEPTH);
+    expect(await this.mock.nextLeafIndex()).to.equal(0n);
+  });
+
+  describe('push', function () {
+    it('tree is correctly updated', async function () {
+      const leafs = Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash);
+
+      // for each leaf slot
+      for (const i in leafs) {
+        // generate random leaf and hash it
+        const hashedLeaf = hashLeaf((leafs[i] = generators.bytes32()));
+
+        // update leaf list and rebuild tree.
+        const tree = makeTree(leafs);
+
+        // push value to tree
+        await expect(this.mock.push(hashedLeaf)).to.emit(this.mock, 'LeafInserted').withArgs(hashedLeaf, i, tree.root);
+
+        // check tree
+        expect(await this.mock.root()).to.equal(tree.root);
+        expect(await this.mock.nextLeafIndex()).to.equal(BigInt(i) + 1n);
+      }
+    });
+
+    it('revert when tree is full', async function () {
+      await Promise.all(Array.from({ length: 2 ** Number(DEPTH) }).map(() => this.mock.push(ethers.ZeroHash)));
+
+      await expect(this.mock.push(ethers.ZeroHash)).to.be.revertedWithPanic(PANIC_CODES.TOO_MUCH_MEMORY_ALLOCATED);
+    });
+  });
+
+  it('reset', async function () {
+    // empty tree
+    const zeroLeafs = Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash);
+    const zeroTree = makeTree(zeroLeafs);
+
+    // tree with one element
+    const leafs = Array.from({ length: 2 ** Number(DEPTH) }, () => ethers.ZeroHash);
+    const hashedLeaf = hashLeaf((leafs[0] = generators.bytes32())); // fill first leaf and hash it
+    const tree = makeTree(leafs);
+
+    // root should be that of a zero tree
+    expect(await this.mock.root()).to.equal(zeroTree.root);
+    expect(await this.mock.nextLeafIndex()).to.equal(0n);
+
+    // push leaf and check root
+    await expect(this.mock.push(hashedLeaf)).to.emit(this.mock, 'LeafInserted').withArgs(hashedLeaf, 0, tree.root);
+
+    expect(await this.mock.root()).to.equal(tree.root);
+    expect(await this.mock.nextLeafIndex()).to.equal(1n);
+
+    // reset tree
+    await this.mock.setup(DEPTH, ZERO);
+
+    expect(await this.mock.root()).to.equal(zeroTree.root);
+    expect(await this.mock.nextLeafIndex()).to.equal(0n);
+
+    // re-push leaf and check root
+    await expect(this.mock.push(hashedLeaf)).to.emit(this.mock, 'LeafInserted').withArgs(hashedLeaf, 0, tree.root);
+
+    expect(await this.mock.root()).to.equal(tree.root);
+    expect(await this.mock.nextLeafIndex()).to.equal(1n);
+  });
+});