@aztec/merkle-tree 0.30.1 → 0.32.0

package/src/snapshots/snapshot_builder_test_suite.ts

@@ -1,14 +1,18 @@
 import { randomBigInt } from '@aztec/foundation/crypto';
+import { Bufferable } from '@aztec/foundation/serialize';
+
+import { jest } from '@jest/globals';
 
 import { TreeBase } from '../tree_base.js';
-import { TreeSnapshotBuilder } from './snapshot_builder.js';
+import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
+
+jest.setTimeout(50_000);
 
 /** Creates a test suit for snapshots */
-export function describeSnapshotBuilderTestSuite<T extends TreeBase, S extends TreeSnapshotBuilder>(
-  getTree: () => T,
-  getSnapshotBuilder: () => S,
-  modifyTree: (tree: T) => Promise<void>,
-) {
+export function describeSnapshotBuilderTestSuite<
+  T extends TreeBase<Bufferable>,
+  S extends TreeSnapshotBuilder<TreeSnapshot<Bufferable>>,
+>(getTree: () => T, getSnapshotBuilder: () => S, modifyTree: (tree: T) => Promise<void>) {
   describe('SnapshotBuilder', () => {
     let tree: T;
     let snapshotBuilder: S;
@@ -34,7 +38,9 @@ export function describeSnapshotBuilderTestSuite<T extends TreeBase, S extends T
 
         const block = 1;
         const snapshot = await snapshotBuilder.snapshot(block);
-        await expect(snapshotBuilder.snapshot(block)).resolves.toEqual(snapshot);
+        const newSnapshot = await snapshotBuilder.snapshot(block);
+
+        expect(newSnapshot.getRoot()).toEqual(snapshot.getRoot());
       });
 
       it('returns the same path if tree has not diverged', async () => {

package/src/sparse_tree/sparse_tree.ts

@@ -1,3 +1,5 @@
+import { Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
+
 import { UpdateOnlyTree } from '../interfaces/update_only_tree.js';
 import { FullTreeSnapshotBuilder } from '../snapshots/full_snapshot.js';
 import { TreeSnapshot } from '../snapshots/snapshot_builder.js';
@@ -6,20 +8,21 @@ import { INITIAL_LEAF, TreeBase } from '../tree_base.js';
 /**
  * A Merkle tree implementation that uses a LevelDB database to store the tree.
  */
-export class SparseTree extends TreeBase implements UpdateOnlyTree {
-  #snapshotBuilder = new FullTreeSnapshotBuilder(this.store, this);
+export class SparseTree<T extends Bufferable> extends TreeBase<T> implements UpdateOnlyTree<T> {
+  #snapshotBuilder = new FullTreeSnapshotBuilder(this.store, this, this.deserializer);
   /**
    * Updates a leaf in the tree.
    * @param leaf - New contents of the leaf.
    * @param index - Index of the leaf to be updated.
    */
-  public updateLeaf(leaf: Buffer, index: bigint): Promise<void> {
+  public updateLeaf(value: T, index: bigint): Promise<void> {
     if (index > this.maxIndex) {
       throw Error(`Index out of bounds. Index ${index}, max index: ${this.maxIndex}.`);
     }
 
+    const leaf = serializeToBuffer(value);
     const insertingZeroElement = leaf.equals(INITIAL_LEAF);
-    const originallyZeroElement = this.getLeafValue(index, true)?.equals(INITIAL_LEAF);
+    const originallyZeroElement = this.getLeafBuffer(index, true)?.equals(INITIAL_LEAF);
     if (insertingZeroElement && originallyZeroElement) {
       return Promise.resolve();
     }
@@ -35,15 +38,19 @@ export class SparseTree extends TreeBase implements UpdateOnlyTree {
     return Promise.resolve();
   }
 
-  public snapshot(block: number): Promise<TreeSnapshot> {
+  public snapshot(block: number): Promise<TreeSnapshot<T>> {
     return this.#snapshotBuilder.snapshot(block);
   }
 
-  public getSnapshot(block: number): Promise<TreeSnapshot> {
+  public getSnapshot(block: number): Promise<TreeSnapshot<T>> {
     return this.#snapshotBuilder.getSnapshot(block);
   }
 
-  public findLeafIndex(_value: Buffer, _includeUncommitted: boolean): bigint | undefined {
+  public findLeafIndex(_value: T, _includeUncommitted: boolean): bigint | undefined {
+    throw new Error('Finding leaf index is not supported for sparse trees');
+  }
+
+  public findLeafIndexAfter(_value: T, _startIndex: bigint, _includeUncommitted: boolean): bigint | undefined {
     throw new Error('Finding leaf index is not supported for sparse trees');
   }
 }

package/src/standard_indexed_tree/standard_indexed_tree.ts

@@ -1,6 +1,7 @@
 import { SiblingPath } from '@aztec/circuit-types';
 import { TreeInsertionStats } from '@aztec/circuit-types/stats';
 import { toBufferBE } from '@aztec/foundation/bigint-buffer';
+import { FromBuffer } from '@aztec/foundation/serialize';
 import { Timer } from '@aztec/foundation/timer';
 import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 import { AztecKVStore, AztecMap } from '@aztec/kv-store';
@@ -51,10 +52,14 @@ function getEmptyLowLeafWitness<N extends number>(
   };
 }
 
+export const noopDeserializer: FromBuffer<Buffer> = {
+  fromBuffer: (buf: Buffer) => buf,
+};
+
 /**
  * Standard implementation of an indexed tree.
  */
-export class StandardIndexedTree extends TreeBase implements IndexedTree {
+export class StandardIndexedTree extends TreeBase<Buffer> implements IndexedTree {
   #snapshotBuilder = new IndexedTreeSnapshotBuilder(this.store, this, this.leafPreimageFactory);
 
   protected cachedLeafPreimages: { [key: string]: IndexedTreeLeafPreimage } = {};
@@ -71,7 +76,7 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree {
     protected leafFactory: LeafFactory,
     root?: Buffer,
   ) {
-    super(store, hasher, name, depth, size, root);
+    super(store, hasher, name, depth, size, noopDeserializer, root);
     this.leaves = store.openMap(`tree_${name}_leaves`);
     this.leafIndex = store.openMap(`tree_${name}_leaf_index`);
   }
@@ -234,6 +239,10 @@ export class StandardIndexedTree extends TreeBase implements IndexedTree {
     return index;
   }
 
+  public findLeafIndexAfter(_leaf: Buffer, _startIndex: bigint, _includeUncommitted: boolean): bigint | undefined {
+    throw new Error('Method not implemented for indexed trees');
+  }
+
   /**
    * Initializes the tree.
    * @param prefilledSize - A number of leaves that are prefilled with values.

package/src/standard_tree/standard_tree.ts

@@ -1,4 +1,5 @@
 import { TreeInsertionStats } from '@aztec/circuit-types/stats';
+import { Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
 import { Timer } from '@aztec/foundation/timer';
 
 import { AppendOnlyTree } from '../interfaces/append_only_tree.js';
@@ -9,15 +10,15 @@ import { TreeBase } from '../tree_base.js';
 /**
  * A Merkle tree implementation that uses a LevelDB database to store the tree.
  */
-export class StandardTree extends TreeBase implements AppendOnlyTree {
-  #snapshotBuilder = new AppendOnlySnapshotBuilder(this.store, this, this.hasher);
+export class StandardTree<T extends Bufferable = Buffer> extends TreeBase<T> implements AppendOnlyTree<T> {
+  #snapshotBuilder = new AppendOnlySnapshotBuilder(this.store, this, this.hasher, this.deserializer);
 
   /**
    * Appends the given leaves to the tree.
    * @param leaves - The leaves to append.
    * @returns Empty promise.
    */
-  public appendLeaves(leaves: Buffer[]): Promise<void> {
+  public appendLeaves(leaves: T[]): Promise<void> {
     this.hasher.reset();
     const timer = new Timer();
     super.appendLeaves(leaves);
@@ -34,18 +35,23 @@ export class StandardTree extends TreeBase implements AppendOnlyTree {
     return Promise.resolve();
   }
 
-  public snapshot(blockNumber: number): Promise<TreeSnapshot> {
+  public snapshot(blockNumber: number): Promise<TreeSnapshot<T>> {
     return this.#snapshotBuilder.snapshot(blockNumber);
   }
 
-  public getSnapshot(blockNumber: number): Promise<TreeSnapshot> {
+  public getSnapshot(blockNumber: number): Promise<TreeSnapshot<T>> {
     return this.#snapshotBuilder.getSnapshot(blockNumber);
   }
 
-  public findLeafIndex(value: Buffer, includeUncommitted: boolean): bigint | undefined {
-    for (let i = 0n; i < this.getNumLeaves(includeUncommitted); i++) {
+  public findLeafIndex(value: T, includeUncommitted: boolean): bigint | undefined {
+    return this.findLeafIndexAfter(value, 0n, includeUncommitted);
+  }
+
+  public findLeafIndexAfter(value: T, startIndex: bigint, includeUncommitted: boolean): bigint | undefined {
+    const buffer = serializeToBuffer(value);
+    for (let i = startIndex; i < this.getNumLeaves(includeUncommitted); i++) {
       const currentValue = this.getLeafValue(i, includeUncommitted);
-      if (currentValue && currentValue.equals(value)) {
+      if (currentValue && serializeToBuffer(currentValue).equals(buffer)) {
         return i;
       }
     }

package/src/tree_base.ts

@@ -1,6 +1,7 @@
 import { SiblingPath } from '@aztec/circuit-types';
 import { toBigIntLE, toBufferLE } from '@aztec/foundation/bigint-buffer';
 import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
+import { Bufferable, FromBuffer, serializeToBuffer } from '@aztec/foundation/serialize';
 import { AztecKVStore, AztecMap, AztecSingleton } from '@aztec/kv-store';
 import { Hasher } from '@aztec/types/interfaces';
 
@@ -44,7 +45,7 @@ export const INITIAL_LEAF = Buffer.from('000000000000000000000000000000000000000
 /**
  * A Merkle tree implementation that uses a LevelDB database to store the tree.
  */
-export abstract class TreeBase implements MerkleTree {
+export abstract class TreeBase<T extends Bufferable> implements MerkleTree<T> {
   protected readonly maxIndex: bigint;
   protected cachedSize?: bigint;
   private root!: Buffer;
@@ -62,6 +63,7 @@ export abstract class TreeBase implements MerkleTree {
     private name: string,
     private depth: number,
     protected size: bigint = 0n,
+    protected deserializer: FromBuffer<T>,
     root?: Buffer,
   ) {
     if (!(depth >= 1 && depth <= MAX_DEPTH)) {
@@ -173,7 +175,22 @@ export abstract class TreeBase implements MerkleTree {
    * @param includeUncommitted - Indicates whether to include uncommitted changes.
    * @returns Leaf value at the given index or undefined.
    */
-  public getLeafValue(index: bigint, includeUncommitted: boolean): Buffer | undefined {
+  public getLeafValue(index: bigint, includeUncommitted: boolean): T | undefined {
+    const buf = this.getLatestValueAtIndex(this.depth, index, includeUncommitted);
+    if (buf) {
+      return this.deserializer.fromBuffer(buf);
+    } else {
+      return undefined;
+    }
+  }
+
+  /**
+   * Gets the value at the given index.
+   * @param index - The index of the leaf.
+   * @param includeUncommitted - Indicates whether to include uncommitted changes.
+   * @returns Leaf value at the given index or undefined.
+   */
+  public getLeafBuffer(index: bigint, includeUncommitted: boolean): Buffer | undefined {
     return this.getLatestValueAtIndex(this.depth, index, includeUncommitted);
   }
 
@@ -292,7 +309,7 @@ export abstract class TreeBase implements MerkleTree {
    *          `getLatestValueAtIndex` will return a value from cache (because at least one of the 2 children was
    *          touched in previous iteration).
    */
-  protected appendLeaves(leaves: Buffer[]): void {
+  protected appendLeaves(leaves: T[]): void {
     const numLeaves = this.getNumLeaves(true);
     if (numLeaves + BigInt(leaves.length) - 1n > this.maxIndex) {
       throw Error(`Can't append beyond max index. Max index: ${this.maxIndex}`);
@@ -303,7 +320,7 @@ export abstract class TreeBase implements MerkleTree {
     let level = this.depth;
     for (let i = 0; i < leaves.length; i++) {
       const cacheKey = indexToKeyHash(this.name, level, firstIndex + BigInt(i));
-      this.cache[cacheKey] = leaves[i];
+      this.cache[cacheKey] = serializeToBuffer(leaves[i]);
     }
 
     let lastIndex = firstIndex + BigInt(leaves.length);
@@ -330,5 +347,14 @@ export abstract class TreeBase implements MerkleTree {
    * @param includeUncommitted - Indicates whether to include uncommitted data.
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  abstract findLeafIndex(value: Buffer, includeUncommitted: boolean): bigint | undefined;
+  abstract findLeafIndex(value: T, includeUncommitted: boolean): bigint | undefined;
+
+  /**
+   * Returns the first index containing a leaf value after `startIndex`.
+   * @param leaf - The leaf value to look for.
+   * @param startIndex - The index to start searching from (used when skipping nullified messages)
+   * @param includeUncommitted - Indicates whether to include uncommitted data.
+   * @returns The index of the first leaf found with a given value (undefined if not found).
+   */
+  abstract findLeafIndexAfter(leaf: T, startIndex: bigint, includeUncommitted: boolean): bigint | undefined;
 }