chore: Compute function tree root in ts.

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