Account Compression TS SDK: add `MerkleTree` static functions (#4048)

account-compression/sdk/package.json

@@ -49,6 +49,7 @@
     "test:events": "start-server-and-test start-validator http://localhost:8899/health run-tests:events",
     "test:accounts": "start-server-and-test start-validator http://localhost:8899/health run-tests:accounts",
     "test:e2e": "start-server-and-test start-validator http://localhost:8899/health run-tests:e2e",
+    "test:merkle-tree": "jest tests/merkleTree.test.ts --detectOpenHandles",
     "test": "start-server-and-test start-validator http://localhost:8899/health run-tests"
   },
   "dependencies": {
@@ -56,6 +57,7 @@
     "@metaplex-foundation/beet-solana": "^0.4.0",
     "bn.js": "^5.2.1",
     "borsh": "^0.7.0",
+    "js-sha3": "^0.8.0",
     "typescript-collections": "^1.3.3"
   },
   "peerDependencies": {
@@ -80,7 +82,6 @@
     "gh-pages": "^4.0.0",
     "jest": "^29.0.1",
     "jest-config": "^29.0.1",
-    "js-sha3": "^0.8.0",
     "start-server-and-test": "^1.14.0",
     "ts-jest": "^28.0.8",
     "ts-jest-resolver": "^2.0.0",

account-compression/sdk/src/merkle-tree/index.ts

@@ -18,6 +18,12 @@ export class MerkleTree {
   root: Buffer;
   depth: number;
 
+  /**
+   * Please use `MerkleTree.sparseMerkleTreeFromLeaves` to
+   * create trees instead. This method is exposed for testing purposes,
+   * and for those that are familiar with the MerkleTree data structure.
+   * @param leaves leaf nodes of the tree
+   */
   constructor(leaves: Buffer[]) {
     let [nodes, finalLeaves] = buildLeaves(leaves);
     let seqNum = leaves.length;
@@ -53,6 +59,30 @@ export class MerkleTree {
     this.depth = nodes.peek()!.level + 1;
   }
 
+  /**
+   * This is the recommended way to create MerkleTrees.
+   * If you're trying to match an on-chain MerkleTree,
+   * set `depth` to `{@link ConcurrentMerkleTreeAccount}.getMaxDepth()`
+   *
+   * @param leaves leaves of the tree
+   * @param depth number of levels in the tree
+   * @returns MerkleTree
+   */
+  static sparseMerkleTreeFromLeaves(
+    leaves: Buffer[],
+    depth: number
+  ): MerkleTree {
+    const _leaves: Buffer[] = [];
+    for (let i = 0; i < 2 ** depth; i++) {
+      if (i < leaves.length) {
+        _leaves.push(leaves[i]);
+      } else {
+        _leaves.push(Buffer.alloc(32));
+      }
+    }
+    return new MerkleTree(_leaves);
+  }
+
   getRoot(): Buffer {
     return this.root;
   }
@@ -127,11 +157,10 @@ export class MerkleTree {
     this.root = node.node;
   }
 
-  verify(
+  static hashProof(
-    root: string,
     merkleTreeProof: MerkleTreeProof,
-    verbose = false
+    verbose: boolean = false
-  ): boolean {
+  ): Buffer {
     const { leaf, leafIndex, proof } = merkleTreeProof;
 
     let node = new PublicKey(leaf).toBuffer();
@@ -143,12 +172,30 @@ export class MerkleTree {
       }
       if (verbose) console.log(`node ${i} ${new PublicKey(node).toString()}`);
     }
+    return node;
+  }
+
+  /**
+   * Verifies that a root matches the proof.
+   * @param root Root of a MerkleTree
+   * @param merkleTreeProof Proof to a leaf in the MerkleTree
+   * @param verbose Whether to print hashed nodes
+   * @returns Whether the proof is valid
+   */
+  static verify(
+    root: Buffer,
+    merkleTreeProof: MerkleTreeProof,
+    verbose: boolean = false
+  ): boolean {
+    const node = MerkleTree.hashProof(merkleTreeProof, verbose);
     const rehashed = new PublicKey(node).toString();
     const received = new PublicKey(root).toString();
-    if (verbose) console.log(`hashed ${rehashed} got ${received}`);
     if (rehashed !== received) {
-      throw new Error("Roots don't match!!!");
+      if (verbose)
+        console.log(`Roots don't match! Expected ${rehashed} got ${received}`);
+      return false;
     }
+    if (verbose) console.log(`Hashed ${rehashed} got ${received}`);
     return rehashed === received;
   }
 }

account-compression/sdk/tests/merkleTree.test.ts

@@ -0,0 +1,34 @@
+import { assert } from "chai";
+import * as crypto from "crypto";
+
+import { emptyNode, MerkleTree } from "../src";
+
+describe("MerkleTree tests", () => {
+  it("Check constructor equivalence for depth 2 tree", () => {
+    const leaves = [
+      crypto.randomBytes(32),
+      crypto.randomBytes(32),
+      crypto.randomBytes(32),
+    ];
+    const rawLeaves = leaves.concat(emptyNode(0));
+    const merkleTreeRaw = new MerkleTree(rawLeaves);
+    const merkleTreeSparse = MerkleTree.sparseMerkleTreeFromLeaves(leaves, 2);
+
+    assert(merkleTreeRaw.root.equals(merkleTreeSparse.root));
+  });
+
+  const TEST_DEPTH = 14;
+  it(`Check proofs for 2^${TEST_DEPTH} tree`, () => {
+    const leaves: Buffer[] = [];
+    for (let i = 0; i < 2 ** TEST_DEPTH; i++) {
+      leaves.push(crypto.randomBytes(32));
+    }
+    const merkleTree = new MerkleTree(leaves);
+
+    // Check proofs
+    for (let i = 0; i < leaves.length; i++) {
+      const proof = merkleTree.getProof(i);
+      assert(MerkleTree.verify(merkleTree.getRoot(), proof));
+    }
+  });
+});