chore: Compute function tree root in ts. (#3326)

One step closer to not needing circuits.wasm.
* This creates a new class `MerkleTreeRootCalculator`. It wouldn't have
had to exist except for the fact that the zero leaf of the function tree
is the hash of a zeroed function leaf (5 fields) rather than just a 0
buffer. Kinda annoying. Needed?
* I tried putting the class in the merkle-tree package, but worryingly
that created a circular dependency because everything depends on this
generic "types" catch all package, and it depends on circuits... Seems
sus. For now the class sits next to where it's used in `abis`.
* When performing a naive tree hash, this actually beats the wasm
performance by around 5%. Perhaps due to the wasm always computing the
zero leaf. Assuming that's the reason, the performance is equal despite
overhead of calling into wasm (so we can assume such overheads are
negligible).
* The algorithm performs much better than the wasm version however, when
the tree is not full, as it leverages zero layer caches to not have to
hash the entire tree. The test adds 4 leaves to a 16 leaf tree, and as
expected is 4 times faster.

Just for clarity, this still uses wasm, but it uses bb.js wasm and at a
lower level. Just calling `pedersenHash`.

yarn-project/circuits.js/src/abis/abis.ts

@@ -11,6 +11,7 @@ import {
   CompleteAddress,
   ContractDeploymentData,
   FUNCTION_SELECTOR_NUM_BYTES,
+  FUNCTION_TREE_HEIGHT,
   Fr,
   FunctionData,
   FunctionLeafPreimage,
@@ -23,7 +24,8 @@ import {
   TxContext,
   TxRequest,
 } from '../index.js';
-import { boolToBuffer, serializeBufferArrayToVector } from '../utils/serialize.js';
+import { boolToBuffer } from '../utils/serialize.js';
+import { MerkleTreeRootCalculator } from './merkle_tree_root_calculator.js';
 
 /**
  * Synchronously calls a wasm function.
@@ -109,16 +111,20 @@ export function computeFunctionLeaf(fnLeaf: FunctionLeafPreimage): Fr {
   );
 }
 
+// The "zero leaf" of the function tree is the hash of 5 zero fields.
+// TODO: Why can we not just use a zero field as the zero leaf? Complicates things perhaps unnecessarily?
+const functionTreeZeroLeaf = pedersenHash(new Array(5).fill(Buffer.alloc(32)));
+const functionTreeRootCalculator = new MerkleTreeRootCalculator(FUNCTION_TREE_HEIGHT, functionTreeZeroLeaf);
+
 /**
  * Computes a function tree root from function leaves.
  * @param wasm - A module providing low-level wasm access.
  * @param fnLeaves - The function leaves to be included in the contract function tree.
  * @returns The function tree root.
  */
 export function computeFunctionTreeRoot(wasm: IWasmModule, fnLeaves: Fr[]) {
-  const inputVector = serializeBufferArrayToVector(fnLeaves.map(fr => fr.toBuffer()));
-  const result = wasmSyncCall(wasm, 'abis__compute_function_tree_root', inputVector, 32);
-  return Fr.fromBuffer(result);
+  const leaves = fnLeaves.map(fr => fr.toBuffer());
+  return Fr.fromBuffer(functionTreeRootCalculator.computeTreeRoot(leaves));
 }
 
 /**

yarn-project/circuits.js/src/abis/merkle_tree_root_calculator.test.ts

@@ -0,0 +1,31 @@
+import { Fr } from '@aztec/foundation/fields';
+
+import { MerkleTreeRootCalculator } from './merkle_tree_root_calculator.js';
+
+describe('merkle tree root calculator', () => {
+  it('should correctly handle no leaves', () => {
+    // Height of 3 is 8 leaves.
+    const calculator = new MerkleTreeRootCalculator(4);
+    const expected = calculator.computeTreeRoot(new Array(8).fill(new Fr(0)).map(fr => fr.toBuffer()));
+    expect(calculator.computeTreeRoot()).toEqual(expected);
+  });
+
+  it('should correctly leverage zero hashes', () => {
+    const calculator = new MerkleTreeRootCalculator(4);
+    const leaves = Array.from({ length: 5 }).map((_, i) => new Fr(i).toBuffer());
+    const padded = [...leaves, ...new Array(3).fill(Buffer.alloc(32))];
+    const expected = calculator.computeTreeRoot(padded);
+    const result = calculator.computeTreeRoot(leaves);
+    expect(result).not.toBeUndefined();
+    expect(result).toEqual(expected);
+  });
+
+  it('should correctly handle non default zero leaf', () => {
+    const zeroLeaf = new Fr(666).toBuffer();
+    const calculator = new MerkleTreeRootCalculator(4, zeroLeaf);
+    const leaves = Array.from({ length: 5 }).map((_, i) => new Fr(i).toBuffer());
+    const padded = [...leaves, ...new Array(3).fill(zeroLeaf)];
+    const expected = calculator.computeTreeRoot(padded);
+    expect(calculator.computeTreeRoot(leaves)).toEqual(expected);
+  });
+});

yarn-project/circuits.js/src/abis/merkle_tree_root_calculator.ts

@@ -0,0 +1,33 @@
+import { pedersenHash } from '@aztec/foundation/crypto';
+
+/**
+ * Calculates the root of a merkle tree.
+ */
+export class MerkleTreeRootCalculator {
+  private zeroHashes: Buffer[];
+
+  constructor(private height: number, zeroLeaf = Buffer.alloc(32)) {
+    this.zeroHashes = Array.from({ length: height }).reduce(
+      (acc: Buffer[], _, i) => [...acc, pedersenHash([acc[i], acc[i]])],
+      [zeroLeaf],
+    );
+  }
+
+  computeTreeRoot(leaves: Buffer[] = []) {
+    if (leaves.length === 0) {
+      return this.zeroHashes[this.zeroHashes.length - 1];
+    }
+
+    for (let i = 0; i < this.height; ++i) {
+      let j = 0;
+      for (; j < leaves.length / 2; ++j) {
+        const l = leaves[j * 2];
+        const r = leaves[j * 2 + 1] || this.zeroHashes[i];
+        leaves[j] = pedersenHash([l, r]);
+      }
+      leaves = leaves.slice(0, j);
+    }
+
+    return leaves[0];
+  }
+}

yarn-project/merkle-tree/package.json

@@ -32,7 +32,6 @@
     "testTimeout": 15000
   },
   "dependencies": {
-    "@aztec/circuits.js": "workspace:^",
     "@aztec/foundation": "workspace:^",
     "@aztec/types": "workspace:^",
     "levelup": "^5.1.1",

yarn-project/merkle-tree/tsconfig.json

@@ -6,9 +6,6 @@
     "tsBuildInfoFile": ".tsbuildinfo"
   },
   "references": [
-    {
-      "path": "../circuits.js"
-    },
     {
       "path": "../foundation"
     },

yarn-project/yarn.lock

@@ -579,7 +579,6 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@aztec/merkle-tree@workspace:merkle-tree"
   dependencies:
-    "@aztec/circuits.js": "workspace:^"
     "@aztec/foundation": "workspace:^"
     "@aztec/types": "workspace:^"
     "@jest/globals": ^29.5.0