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]>

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;
 }

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

@@ -1,14 +1,18 @@
-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 updates at arbitrary indices but not appending.
  */
-export interface UpdateOnlyTree extends MerkleTree, TreeSnapshotBuilder {
+export interface UpdateOnlyTree<T extends Bufferable = Buffer>
+  extends MerkleTree<T>,
+    TreeSnapshotBuilder<TreeSnapshot<T>> {
   /**
    * Updates a leaf at a given index in the tree.
    * @param leaf - The leaf value to be updated.
    * @param index - The leaf to be updated.
    */
-  updateLeaf(leaf: Buffer, index: bigint): Promise<void>;
+  updateLeaf(leaf: T, index: bigint): Promise<void>;
 }

yarn-project/merkle-tree/src/load_tree.ts

@@ -1,3 +1,4 @@
+import { Bufferable, FromBuffer } from '@aztec/foundation/serialize';
 import { AztecKVStore } from '@aztec/kv-store';
 import { Hasher } from '@aztec/types/interfaces';
 
@@ -11,13 +12,22 @@ import { TreeBase, getTreeMeta } from './tree_base.js';
  * @param name - Name of the tree.
  * @returns The newly created tree.
  */
-export function loadTree<T extends TreeBase>(
-  c: new (store: AztecKVStore, hasher: Hasher, name: string, depth: number, size: bigint, root: Buffer) => T,
+export function loadTree<T extends TreeBase<Bufferable>, D extends FromBuffer<Bufferable>>(
+  c: new (
+    store: AztecKVStore,
+    hasher: Hasher,
+    name: string,
+    depth: number,
+    size: bigint,
+    deserializer: D,
+    root: Buffer,
+  ) => T,
   store: AztecKVStore,
   hasher: Hasher,
   name: string,
+  deserializer: D,
 ): Promise<T> {
   const { root, depth, size } = getTreeMeta(store, name);
-  const tree = new c(store, hasher, name, depth, size, root);
+  const tree = new c(store, hasher, name, depth, size, deserializer, root);
   return Promise.resolve(tree);
 }

yarn-project/merkle-tree/src/new_tree.ts

@@ -1,3 +1,4 @@
+import { Bufferable, FromBuffer } from '@aztec/foundation/serialize';
 import { AztecKVStore } from '@aztec/kv-store';
 import { Hasher } from '@aztec/types/interfaces';
 
@@ -13,15 +14,16 @@ import { TreeBase } from './tree_base.js';
  * @param prefilledSize - A number of leaves that are prefilled with values.
  * @returns The newly created tree.
  */
-export async function newTree<T extends TreeBase>(
-  c: new (store: AztecKVStore, hasher: Hasher, name: string, depth: number, size: bigint) => T,
+export async function newTree<T extends TreeBase<Bufferable>, D extends FromBuffer<Bufferable>>(
+  c: new (store: AztecKVStore, hasher: Hasher, name: string, depth: number, size: bigint, deserializer: D) => T,
   store: AztecKVStore,
   hasher: Hasher,
   name: string,
+  deserializer: D,
   depth: number,
   prefilledSize = 1,
 ): Promise<T> {
-  const tree = new c(store, hasher, name, depth, 0n);
+  const tree = new c(store, hasher, name, depth, 0n, deserializer);
   await tree.init(prefilledSize);
   return tree;
 }

yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts

@@ -1,4 +1,5 @@
 import { randomBytes } from '@aztec/foundation/crypto';
+import { FromBuffer } from '@aztec/foundation/serialize';
 import { AztecKVStore } from '@aztec/kv-store';
 import { openTmpStore } from '@aztec/kv-store/utils';
 
@@ -8,14 +9,15 @@ import { describeSnapshotBuilderTestSuite } from './snapshot_builder_test_suite.
 
 describe('AppendOnlySnapshot', () => {
   let tree: StandardTree;
-  let snapshotBuilder: AppendOnlySnapshotBuilder;
+  let snapshotBuilder: AppendOnlySnapshotBuilder<Buffer>;
   let db: AztecKVStore;
 
   beforeEach(async () => {
     db = openTmpStore();
     const hasher = new Pedersen();
-    tree = await newTree(StandardTree, db, hasher, 'test', 4);
-    snapshotBuilder = new AppendOnlySnapshotBuilder(db, tree, hasher);
+    const deserializer: FromBuffer<Buffer> = { fromBuffer: b => b };
+    tree = await newTree(StandardTree, db, hasher, 'test', deserializer, 4);
+    snapshotBuilder = new AppendOnlySnapshotBuilder(db, tree, hasher, deserializer);
   });
 
   describeSnapshotBuilderTestSuite(