@aztec/merkle-tree 0.31.0 → 0.32.1

package/src/snapshots/full_snapshot.ts

@@ -1,6 +1,9 @@
-import { TreeBase } from '../tree_base.js';
+import { type Bufferable, type FromBuffer } from '@aztec/foundation/serialize';
+import { type AztecKVStore } from '@aztec/kv-store';
+
+import { type TreeBase } from '../tree_base.js';
 import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js';
-import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
+import { type TreeSnapshot, type TreeSnapshotBuilder } from './snapshot_builder.js';
 
 /**
  * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores
@@ -16,11 +19,15 @@ import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
  * Worst case space complexity: O(N * M)
  * Sibling path access: O(H) database reads
  */
-export class FullTreeSnapshotBuilder
-  extends BaseFullTreeSnapshotBuilder<TreeBase, TreeSnapshot>
-  implements TreeSnapshotBuilder<TreeSnapshot>
+export class FullTreeSnapshotBuilder<T extends Bufferable>
+  extends BaseFullTreeSnapshotBuilder<TreeBase<T>, TreeSnapshot<T>>
+  implements TreeSnapshotBuilder<TreeSnapshot<T>>
 {
-  protected openSnapshot(root: Buffer, numLeaves: bigint): TreeSnapshot {
-    return new BaseFullTreeSnapshot(this.nodes, root, numLeaves, this.tree);
+  constructor(db: AztecKVStore, tree: TreeBase<T>, private deserializer: FromBuffer<T>) {
+    super(db, tree);
+  }
+
+  protected openSnapshot(root: Buffer, numLeaves: bigint): TreeSnapshot<T> {
+    return new BaseFullTreeSnapshot(this.nodes, root, numLeaves, this.tree, this.deserializer);
   }
 }

package/src/snapshots/indexed_tree_snapshot.ts

@@ -1,20 +1,20 @@
-import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
-import { AztecKVStore, AztecMap } from '@aztec/kv-store';
+import { type IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
 
-import { IndexedTree, PreimageFactory } from '../interfaces/indexed_tree.js';
-import { TreeBase } from '../tree_base.js';
+import { type IndexedTree, type PreimageFactory } from '../interfaces/indexed_tree.js';
+import { type TreeBase } from '../tree_base.js';
 import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js';
-import { IndexedTreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
+import { type IndexedTreeSnapshot, type TreeSnapshotBuilder } from './snapshot_builder.js';
 
 const snapshotLeafValue = (node: Buffer, index: bigint) => 'snapshot:leaf:' + node.toString('hex') + ':' + index;
 
 /** a */
 export class IndexedTreeSnapshotBuilder
-  extends BaseFullTreeSnapshotBuilder<IndexedTree & TreeBase, IndexedTreeSnapshot>
+  extends BaseFullTreeSnapshotBuilder<IndexedTree & TreeBase<Buffer>, IndexedTreeSnapshot>
   implements TreeSnapshotBuilder<IndexedTreeSnapshot>
 {
   leaves: AztecMap<string, Buffer>;
-  constructor(store: AztecKVStore, tree: IndexedTree & TreeBase, private leafPreimageBuilder: PreimageFactory) {
+  constructor(store: AztecKVStore, tree: IndexedTree & TreeBase<Buffer>, private leafPreimageBuilder: PreimageFactory) {
     super(store, tree);
     this.leaves = store.openMap('indexed_tree_snapshot:' + tree.getName());
   }
@@ -32,16 +32,16 @@ export class IndexedTreeSnapshotBuilder
 }
 
 /** A snapshot of an indexed tree at a particular point in time */
-class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTreeSnapshot {
+class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot<Buffer> implements IndexedTreeSnapshot {
   constructor(
     db: AztecMap<string, [Buffer, Buffer]>,
     private leaves: AztecMap<string, Buffer>,
     historicRoot: Buffer,
     numLeaves: bigint,
-    tree: IndexedTree & TreeBase,
+    tree: IndexedTree & TreeBase<Buffer>,
     private leafPreimageBuilder: PreimageFactory,
   ) {
-    super(db, historicRoot, numLeaves, tree);
+    super(db, historicRoot, numLeaves, tree, { fromBuffer: buf => buf });
   }
 
   getLeafValue(index: bigint): Buffer | undefined {

package/src/snapshots/snapshot_builder.ts

@@ -1,10 +1,11 @@
-import { SiblingPath } from '@aztec/circuit-types';
-import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { type SiblingPath } from '@aztec/circuit-types';
+import { type Bufferable } from '@aztec/foundation/serialize';
+import { type IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 
 /**
  * An interface for a tree that can record snapshots of its contents.
  */
-export interface TreeSnapshotBuilder<S extends TreeSnapshot = TreeSnapshot> {
+export interface TreeSnapshotBuilder<S extends TreeSnapshot<Bufferable>> {
   /**
    * Creates a snapshot of the tree at the given version.
    * @param block - The version to snapshot the tree at.
@@ -21,7 +22,7 @@ export interface TreeSnapshotBuilder<S extends TreeSnapshot = TreeSnapshot> {
 /**
  * A tree snapshot
  */
-export interface TreeSnapshot {
+export interface TreeSnapshot<T extends Bufferable> {
   /**
    * Returns the current root of the tree.
    */
@@ -41,7 +42,7 @@ export interface TreeSnapshot {
    * Returns the value of a leaf at the specified index.
    * @param index - The index of the leaf value to be returned.
    */
-  getLeafValue(index: bigint): Buffer | undefined;
+  getLeafValue(index: bigint): T | undefined;
 
   /**
    * Returns the sibling path for a requested leaf index.
@@ -55,7 +56,7 @@ export interface TreeSnapshot {
    * @param value - The leaf value to look for.
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  findLeafIndex(value: Buffer): bigint | undefined;
+  findLeafIndex(value: T): bigint | undefined;
 
   /**
    * Returns the first index containing a leaf value after `startIndex`.
@@ -63,11 +64,11 @@ export interface TreeSnapshot {
    * @param startIndex - The index to start searching from (used when skipping nullified messages)
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  findLeafIndexAfter(leaf: Buffer, startIndex: bigint): bigint | undefined;
+  findLeafIndexAfter(leaf: T, startIndex: bigint): bigint | undefined;
 }
 
 /** A snapshot of an indexed tree */
-export interface IndexedTreeSnapshot extends TreeSnapshot {
+export interface IndexedTreeSnapshot extends TreeSnapshot<Buffer> {
   /**
    * Gets the historical data for a leaf
    * @param index - The index of the leaf to get the data for

package/src/snapshots/snapshot_builder_test_suite.ts

@@ -1,14 +1,18 @@
 import { randomBigInt } from '@aztec/foundation/crypto';
+import { type Bufferable } from '@aztec/foundation/serialize';
 
-import { TreeBase } from '../tree_base.js';
-import { TreeSnapshotBuilder } from './snapshot_builder.js';
+import { jest } from '@jest/globals';
+
+import { type TreeBase } from '../tree_base.js';
+import { type TreeSnapshot, type 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,25 +1,28 @@
-import { UpdateOnlyTree } from '../interfaces/update_only_tree.js';
+import { type Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
+
+import { type UpdateOnlyTree } from '../interfaces/update_only_tree.js';
 import { FullTreeSnapshotBuilder } from '../snapshots/full_snapshot.js';
-import { TreeSnapshot } from '../snapshots/snapshot_builder.js';
+import { type TreeSnapshot } from '../snapshots/snapshot_builder.js';
 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,19 +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: Buffer, _startIndex: bigint, _includeUncommitted: boolean): bigint | undefined {
+  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,14 +1,20 @@
 import { SiblingPath } from '@aztec/circuit-types';
-import { TreeInsertionStats } from '@aztec/circuit-types/stats';
+import { type TreeInsertionStats } from '@aztec/circuit-types/stats';
 import { toBufferBE } from '@aztec/foundation/bigint-buffer';
+import { type 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';
-import { Hasher } from '@aztec/types/interfaces';
-
-import { BatchInsertionResult, IndexedTree, LowLeafWitnessData, PreimageFactory } from '../interfaces/indexed_tree.js';
+import { type IndexedTreeLeaf, type IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
+import { type Hasher } from '@aztec/types/interfaces';
+
+import {
+  type BatchInsertionResult,
+  type IndexedTree,
+  type LowLeafWitnessData,
+  type PreimageFactory,
+} from '../interfaces/indexed_tree.js';
 import { IndexedTreeSnapshotBuilder } from '../snapshots/indexed_tree_snapshot.js';
-import { IndexedTreeSnapshot } from '../snapshots/snapshot_builder.js';
+import { type IndexedTreeSnapshot } from '../snapshots/snapshot_builder.js';
 import { TreeBase } from '../tree_base.js';
 
 export const buildDbKeyForPreimage = (name: string, index: bigint) => {
@@ -51,10 +57,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 +81,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`);
   }

package/src/standard_tree/standard_tree.ts

@@ -1,23 +1,24 @@
-import { TreeInsertionStats } from '@aztec/circuit-types/stats';
+import { type TreeInsertionStats } from '@aztec/circuit-types/stats';
+import { type Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
 import { Timer } from '@aztec/foundation/timer';
 
-import { AppendOnlyTree } from '../interfaces/append_only_tree.js';
+import { type AppendOnlyTree } from '../interfaces/append_only_tree.js';
 import { AppendOnlySnapshotBuilder } from '../snapshots/append_only_snapshot.js';
-import { TreeSnapshot } from '../snapshots/snapshot_builder.js';
+import { type TreeSnapshot } from '../snapshots/snapshot_builder.js';
 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,22 +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 {
+  public findLeafIndex(value: T, includeUncommitted: boolean): bigint | undefined {
     return this.findLeafIndexAfter(value, 0n, includeUncommitted);
   }
 
-  public findLeafIndexAfter(value: Buffer, startIndex: bigint, includeUncommitted: boolean): bigint | undefined {
+  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,11 +1,12 @@
 import { SiblingPath } from '@aztec/circuit-types';
 import { toBigIntLE, toBufferLE } from '@aztec/foundation/bigint-buffer';
-import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
-import { AztecKVStore, AztecMap, AztecSingleton } from '@aztec/kv-store';
-import { Hasher } from '@aztec/types/interfaces';
+import { type DebugLogger, createDebugLogger } from '@aztec/foundation/log';
+import { type Bufferable, type FromBuffer, serializeToBuffer } from '@aztec/foundation/serialize';
+import { type AztecKVStore, type AztecMap, type AztecSingleton } from '@aztec/kv-store';
+import { type Hasher } from '@aztec/types/interfaces';
 
 import { HasherWithStats } from './hasher_with_stats.js';
-import { MerkleTree } from './interfaces/merkle_tree.js';
+import { type MerkleTree } from './interfaces/merkle_tree.js';
 
 const MAX_DEPTH = 254;
 
@@ -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,7 +347,7 @@ 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`.
@@ -339,5 +356,5 @@ 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 findLeafIndexAfter(leaf: Buffer, startIndex: bigint, includeUncommitted: boolean): bigint | undefined;
+  abstract findLeafIndexAfter(leaf: T, startIndex: bigint, includeUncommitted: boolean): bigint | undefined;
 }