refactor: append-only merkle tree generics (#5355)

Give Merkle Trees generic type arguments for the leaves they store.

---------

Co-authored-by: Santiago Palladino <[email protected]>

Author: alexghr

cspell.json

@@ -79,6 +79,7 @@
     "erc",
     "falsey",
     "fargate",
+    "Fieldeable",
     "filestat",
     "flatmap",
     "foundryup",

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -458,8 +458,15 @@ export class AztecNodeService implements AztecNode {
 
     const treeHeight = Math.ceil(Math.log2(l2ToL1Messages.length));
     // The root of this tree is the out_hash calculated in Noir => we truncate to match Noir's SHA
-    const tree = new StandardTree(openTmpStore(true), new SHA256Trunc(), 'temp_outhash_sibling_path', treeHeight);
-    await tree.appendLeaves(l2ToL1Messages.map(l2ToL1Msg => l2ToL1Msg.toBuffer()));
+    const tree = new StandardTree(
+      openTmpStore(true),
+      new SHA256Trunc(),
+      'temp_outhash_sibling_path',
+      treeHeight,
+      0n,
+      Fr,
+    );
+    await tree.appendLeaves(l2ToL1Messages);
 
     return [indexOfL2ToL1Message, await tree.getSiblingPath(indexOfL2ToL1Message, true)];
   }

yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts

@@ -38,14 +38,14 @@ describe('e2e_blacklist_token_contract', () => {
 
   let tokenSim: TokenSimulator;
 
-  let slowUpdateTreeSimulator: SparseTree;
+  let slowUpdateTreeSimulator: SparseTree<Fr>;
 
   let cheatCodes: CheatCodes;
 
   const getMembershipProof = async (index: bigint, includeUncommitted: boolean) => {
     return {
       index,
-      value: Fr.fromBuffer(slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted)!),
+      value: slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted)!,
       // eslint-disable-next-line camelcase
       sibling_path: (await slowUpdateTreeSimulator.getSiblingPath(index, includeUncommitted)).toFields(),
     };
@@ -101,9 +101,9 @@ describe('e2e_blacklist_token_contract', () => {
     await wallets[accountIndex].addNote(extendedNote);
   };
 
-  const updateSlowTree = async (tree: SparseTree, wallet: Wallet, index: AztecAddress, value: bigint) => {
+  const updateSlowTree = async (tree: SparseTree<Fr>, wallet: Wallet, index: AztecAddress, value: bigint) => {
     await wallet.addCapsule(getUpdateCapsule(await getUpdateProof(value, index.toBigInt())));
-    await tree.updateLeaf(new Fr(value).toBuffer(), index.toBigInt());
+    await tree.updateLeaf(new Fr(value), index.toBigInt());
   };
 
   beforeAll(async () => {
@@ -113,7 +113,7 @@ describe('e2e_blacklist_token_contract', () => {
     slowTree = await SlowTreeContract.deploy(wallets[0]).send().deployed();
 
     const depth = 254;
-    slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', depth);
+    slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', Fr, depth);
 
     // Add account[0] as admin
     await updateSlowTree(slowUpdateTreeSimulator, wallets[0], accounts[0].address, 4n);

yarn-project/end-to-end/src/e2e_dapp_subscription.test.ts

@@ -28,7 +28,7 @@ import {
   setup,
 } from './fixtures/utils.js';
 
-jest.setTimeout(1_000_000);
+jest.setTimeout(100_000);
 
 const TOKEN_NAME = 'BananaCoin';
 const TOKEN_SYMBOL = 'BC';

yarn-project/end-to-end/src/e2e_slow_tree.test.ts

@@ -23,11 +23,11 @@ describe('e2e_slow_tree', () => {
 
   it('Messing around with noir slow tree', async () => {
     const depth = 254;
-    const slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', depth);
+    const slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', Fr, depth);
     const getMembershipProof = async (index: bigint, includeUncommitted: boolean) => {
       return {
         index,
-        value: Fr.fromBuffer(slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted)!),
+        value: slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted)!,
         // eslint-disable-next-line camelcase
         sibling_path: (await slowUpdateTreeSimulator.getSiblingPath(index, includeUncommitted)).toFields(),
       };
@@ -102,7 +102,7 @@ describe('e2e_slow_tree', () => {
       .update_at_public(await getUpdateProof(1n, key))
       .send()
       .wait();
-    await slowUpdateTreeSimulator.updateLeaf(new Fr(1).toBuffer(), key);
+    await slowUpdateTreeSimulator.updateLeaf(new Fr(1), key);
 
     // Update below.
     _root = {
@@ -140,7 +140,7 @@ describe('e2e_slow_tree', () => {
     const t2 = computeNextChange(BigInt(await cheatCodes.eth.timestamp()));
     await wallet.addCapsule(getUpdateCapsule(await getUpdateProof(4n, key)));
     await contract.methods.update_at_private(key, 4n).send().wait();
-    await slowUpdateTreeSimulator.updateLeaf(new Fr(4).toBuffer(), key);
+    await slowUpdateTreeSimulator.updateLeaf(new Fr(4), key);
     _root = {
       before: Fr.fromBuffer(slowUpdateTreeSimulator.getRoot(false)).toBigInt(),
       after: Fr.fromBuffer(slowUpdateTreeSimulator.getRoot(true)).toBigInt(),

yarn-project/end-to-end/src/integration_l1_publisher.test.ts

@@ -428,8 +428,15 @@ describe('L1Publisher integration', () => {
 
       const treeHeight = Math.ceil(Math.log2(newL2ToL1MsgsArray.length));
 
-      const tree = new StandardTree(openTmpStore(true), new SHA256Trunc(), 'temp_outhash_sibling_path', treeHeight);
-      await tree.appendLeaves(newL2ToL1MsgsArray.map(l2ToL1Msg => l2ToL1Msg.toBuffer()));
+      const tree = new StandardTree(
+        openTmpStore(true),
+        new SHA256Trunc(),
+        'temp_outhash_sibling_path',
+        treeHeight,
+        0n,
+        Fr,
+      );
+      await tree.appendLeaves(newL2ToL1MsgsArray);
 
       const expectedRoot = tree.getRoot(true);
       const [actualRoot] = await outbox.read.roots([block.header.globalVariables.blockNumber.toBigInt()]);

yarn-project/foundation/src/serialize/buffer_reader.ts

@@ -312,3 +312,14 @@ export class BufferReader {
     return this.buffer.length;
   }
 }
+
+/**
+ * A deserializer
+ */
+export interface FromBuffer<T> {
+  /**
+   * Deserializes an object from a buffer
+   * @param buffer - The buffer to deserialize.
+   */
+  fromBuffer(buffer: Buffer): T;
+}

yarn-project/merkle-tree/src/interfaces/append_only_tree.ts

@@ -1,13 +1,17 @@
-import { TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js';
+import { Bufferable } from '@aztec/foundation/serialize';
+
+import { TreeSnapshot, TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js';
 import { MerkleTree } from './merkle_tree.js';
 
 /**
  * A Merkle tree that supports only appending leaves and not updating existing leaves.
  */
-export interface AppendOnlyTree extends MerkleTree, TreeSnapshotBuilder {
+export interface AppendOnlyTree<T extends Bufferable = Buffer>
+  extends MerkleTree<T>,
+    TreeSnapshotBuilder<TreeSnapshot<T>> {
   /**
    * Appends a set of leaf values to the tree.
    * @param leaves - The set of leaves to be appended.
    */
-  appendLeaves(leaves: Buffer[]): Promise<void>;
+  appendLeaves(leaves: T[]): Promise<void>;
 }

yarn-project/merkle-tree/src/interfaces/indexed_tree.ts

@@ -1,7 +1,9 @@
 import { SiblingPath } from '@aztec/circuit-types';
 import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 
+import { IndexedTreeSnapshot, TreeSnapshot, TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js';
 import { AppendOnlyTree } from './append_only_tree.js';
+import { MerkleTree } from './merkle_tree.js';
 
 /**
  * Factory for creating leaf preimages.
@@ -73,7 +75,10 @@ export interface BatchInsertionResult<TreeHeight extends number, SubtreeSiblingP
 /**
  * Indexed merkle tree.
  */
-export interface IndexedTree extends AppendOnlyTree {
+export interface IndexedTree
+  extends MerkleTree<Buffer>,
+    TreeSnapshotBuilder<IndexedTreeSnapshot>,
+    Omit<AppendOnlyTree<Buffer>, keyof TreeSnapshotBuilder<TreeSnapshot<Buffer>>> {
   /**
    * Finds the index of the largest leaf whose value is less than or equal to the provided value.
    * @param newValue - The new value to be inserted into the tree.

yarn-project/merkle-tree/src/interfaces/merkle_tree.ts

@@ -1,4 +1,5 @@
 import { SiblingPath } from '@aztec/circuit-types';
+import { Bufferable } from '@aztec/foundation/serialize';
 
 /**
  * Defines the interface for a source of sibling paths.
@@ -15,7 +16,7 @@ export interface SiblingPathSource {
 /**
  * Defines the interface for a Merkle tree.
  */
-export interface MerkleTree extends SiblingPathSource {
+export interface MerkleTree<T extends Bufferable = Buffer> extends SiblingPathSource {
   /**
    * Returns the current root of the tree.
    * @param includeUncommitted - Set to true to include uncommitted updates in the calculated root.
@@ -48,15 +49,15 @@ export interface MerkleTree extends SiblingPathSource {
    * @param index - The index of the leaf value to be returned.
    * @param includeUncommitted - Set to true to include uncommitted updates in the data set.
    */
-  getLeafValue(index: bigint, includeUncommitted: boolean): Buffer | undefined;
+  getLeafValue(index: bigint, includeUncommitted: boolean): T | undefined;
 
   /**
    * Returns the index of a leaf given its value, or undefined if no leaf with that value is found.
    * @param leaf - The leaf value to look for.
    * @param includeUncommitted - Indicates whether to include uncommitted data.
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  findLeafIndex(leaf: Buffer, includeUncommitted: boolean): bigint | undefined;
+  findLeafIndex(leaf: T, includeUncommitted: boolean): bigint | undefined;
 
   /**
    * Returns the first index containing a leaf value after `startIndex`.
@@ -65,5 +66,5 @@ export interface MerkleTree extends SiblingPathSource {
    * @param includeUncommitted - Indicates whether to include uncommitted data.
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  findLeafIndexAfter(leaf: Buffer, startIndex: bigint, includeUncommitted: boolean): bigint | undefined;
+  findLeafIndexAfter(leaf: T, startIndex: bigint, includeUncommitted: boolean): bigint | undefined;
 }