feat: refactor StandardIndexedTree for abstract leaves and preimages and optimized it

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

@@ -8,6 +8,7 @@ import {
   L1_TO_L2_MSG_TREE_HEIGHT,
   NOTE_HASH_TREE_HEIGHT,
   NULLIFIER_TREE_HEIGHT,
+  NullifierLeafPreimage,
   PUBLIC_DATA_TREE_HEIGHT,
 } from '@aztec/circuits.js';
 import { computeGlobalsHash, computePublicDataTreeIndex } from '@aztec/circuits.js/abis';
@@ -429,19 +430,19 @@ export class AztecNodeService implements AztecNode {
       return undefined;
     }
 
-    const leafDataPromise = db.getLeafData(MerkleTreeId.NULLIFIER_TREE, Number(index));
+    const leafPreimagePromise = db.getLeafPreimage(MerkleTreeId.NULLIFIER_TREE, index);
     const siblingPathPromise = db.getSiblingPath<typeof NULLIFIER_TREE_HEIGHT>(
       MerkleTreeId.NULLIFIER_TREE,
       BigInt(index),
     );
 
-    const [leafData, siblingPath] = await Promise.all([leafDataPromise, siblingPathPromise]);
+    const [leafPreimage, siblingPath] = await Promise.all([leafPreimagePromise, siblingPathPromise]);
 
-    if (!leafData) {
+    if (!leafPreimage) {
       return undefined;
     }
 
-    return new NullifierMembershipWitness(BigInt(index), leafData, siblingPath);
+    return new NullifierMembershipWitness(BigInt(index), leafPreimage as NullifierLeafPreimage, siblingPath);
   }
 
   /**
@@ -463,22 +464,21 @@ export class AztecNodeService implements AztecNode {
     nullifier: Fr,
   ): Promise<NullifierMembershipWitness | undefined> {
     const committedDb = await this.#getWorldState(blockNumber);
-    const { index, alreadyPresent } = await committedDb.getPreviousValueIndex(
-      MerkleTreeId.NULLIFIER_TREE,
-      nullifier.toBigInt(),
-    );
+    const findResult = await committedDb.getPreviousValueIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBigInt());
+    if (!findResult) {
+      return undefined;
+    }
+    const { index, alreadyPresent } = findResult;
     if (alreadyPresent) {
       this.log.warn(`Nullifier ${nullifier.toBigInt()} already exists in the tree`);
     }
-    const leafData = await committedDb.getLeafData(MerkleTreeId.NULLIFIER_TREE, index);
-    if (!leafData) {
-      return undefined;
-    }
+    const preimageData = (await committedDb.getLeafPreimage(MerkleTreeId.NULLIFIER_TREE, index))!;
+
     const siblingPath = await committedDb.getSiblingPath<typeof NULLIFIER_TREE_HEIGHT>(
       MerkleTreeId.NULLIFIER_TREE,
       BigInt(index),
     );
-    return new NullifierMembershipWitness(BigInt(index), leafData, siblingPath);
+    return new NullifierMembershipWitness(BigInt(index), preimageData as NullifierLeafPreimage, siblingPath);
   }
 
   /**

yarn-project/circuits.js/src/structs/rollup/base_rollup.ts

@@ -1,5 +1,7 @@
+import { toBigIntBE, toBufferBE } from '@aztec/foundation/bigint-buffer';
 import { Fr } from '@aztec/foundation/fields';
 import { BufferReader, Tuple } from '@aztec/foundation/serialize';
+import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 
 import {
   BLOCKS_TREE_HEIGHT,
@@ -25,28 +27,103 @@ import { AppendOnlyTreeSnapshot } from './append_only_tree_snapshot.js';
  * Class containing the data of a preimage of a single leaf in the nullifier tree.
  * Note: It's called preimage because this data gets hashed before being inserted as a node into the `IndexedTree`.
  */
-export class NullifierLeafPreimage {
+export class NullifierLeafPreimage implements IndexedTreeLeafPreimage {
   constructor(
     /**
      * Leaf value inside the indexed tree's linked list.
      */
-    public leafValue: Fr,
+    public nullifier: Fr,
     /**
      * Next value inside the indexed tree's linked list.
      */
-    public nextValue: Fr,
+    public nextNullifier: Fr,
     /**
      * Index of the next leaf in the indexed tree's linked list.
      */
-    public nextIndex: UInt32,
+    public nextIndex: bigint,
   ) {}
 
-  toBuffer() {
-    return serializeToBuffer(this.leafValue, this.nextValue, this.nextIndex);
+  getKey(): bigint {
+    return this.nullifier.toBigInt();
+  }
+
+  getNextKey(): bigint {
+    return this.nextNullifier.toBigInt();
+  }
+
+  getNextIndex(): bigint {
+    return this.nextIndex;
+  }
+
+  asLeaf(): NullifierLeaf {
+    return new NullifierLeaf(this.nullifier);
+  }
+
+  toBuffer(): Buffer {
+    return Buffer.concat(this.toHashInputs());
+  }
+
+  toHashInputs(): Buffer[] {
+    return [
+      Buffer.from(this.nullifier.toBuffer()),
+      Buffer.from(toBufferBE(this.nextIndex, 32)),
+      Buffer.from(this.nextNullifier.toBuffer()),
+    ];
+  }
+
+  clone(): NullifierLeafPreimage {
+    return new NullifierLeafPreimage(this.nullifier, this.nextNullifier, this.nextIndex);
+  }
+
+  static empty(): NullifierLeafPreimage {
+    return new NullifierLeafPreimage(Fr.ZERO, Fr.ZERO, 0n);
+  }
+
+  static fromBuffer(buf: Buffer): NullifierLeafPreimage {
+    const nullifier = Fr.fromBuffer(buf.subarray(0, 32));
+    const nextIndex = toBigIntBE(buf.subarray(32, 64));
+    const nextNullifier = Fr.fromBuffer(buf.subarray(64, 96));
+    return new NullifierLeafPreimage(nullifier, nextNullifier, nextIndex);
+  }
+
+  static fromLeaf(leaf: NullifierLeaf, nextKey: bigint, nextIndex: bigint): NullifierLeafPreimage {
+    return new NullifierLeafPreimage(leaf.nullifier, new Fr(nextKey), nextIndex);
+  }
+
+  static clone(preimage: NullifierLeafPreimage): NullifierLeafPreimage {
+    return new NullifierLeafPreimage(preimage.nullifier, preimage.nextNullifier, preimage.nextIndex);
+  }
+}
+
+/**
+ * A nullifier to be inserted in the nullifier tree.
+ */
+export class NullifierLeaf implements IndexedTreeLeaf {
+  constructor(
+    /**
+     * Nullifier value.
+     */
+    public nullifier: Fr,
+  ) {}
+
+  getKey(): bigint {
+    return this.nullifier.toBigInt();
+  }
+
+  toBuffer(): Buffer {
+    return this.nullifier.toBuffer();
+  }
+
+  isEmpty(): boolean {
+    return this.nullifier.isZero();
+  }
+
+  static buildDummy(key: bigint): NullifierLeaf {
+    return new NullifierLeaf(new Fr(key));
   }
 
-  static empty() {
-    return new NullifierLeafPreimage(Fr.ZERO, Fr.ZERO, 0);
+  static fromBuffer(buf: Buffer): NullifierLeaf {
+    return new NullifierLeaf(Fr.fromBuffer(buf));
   }
 }
 

yarn-project/circuits.js/src/tests/factories.ts

@@ -900,7 +900,7 @@ export function makeBaseRollupInputs(seed = 0): BaseRollupInputs {
 
   const lowNullifierLeafPreimages = makeTuple(
     MAX_NEW_NULLIFIERS_PER_BASE_ROLLUP,
-    x => new NullifierLeafPreimage(fr(x), fr(x + 0x100), x + 0x200),
+    x => new NullifierLeafPreimage(fr(x), fr(x + 0x100), BigInt(x + 0x200)),
     seed + 0x1000,
   );
 

yarn-project/end-to-end/tsconfig.json

@@ -55,7 +55,5 @@
       "path": "../world-state"
     }
   ],
-  "include": [
-    "src"
-  ]
+  "include": ["src"]
 }

yarn-project/foundation/package.json

@@ -28,6 +28,7 @@
     "./sleep": "./dest/sleep/index.js",
     "./timer": "./dest/timer/index.js",
     "./transport": "./dest/transport/index.js",
+    "./trees": "./dest/trees/index.js",
     "./wasm": "./dest/wasm/index.js",
     "./worker": "./dest/worker/index.js",
     "./bigint-buffer": "./dest/bigint-buffer/index.js",

yarn-project/foundation/src/index.ts

@@ -21,6 +21,7 @@ export * as serialize from './serialize/index.js';
 export * as sleep from './sleep/index.js';
 export * as timer from './timer/index.js';
 export * as transport from './transport/index.js';
+export * as trees from './trees/index.js';
 export * as types from './types/index.js';
 export * as url from './url/index.js';
 export * as wasm from './wasm/index.js';

yarn-project/foundation/src/trees/index.ts

@@ -0,0 +1,48 @@
+/**
+ * A leaf of an indexed merkle tree.
+ */
+export interface IndexedTreeLeaf {
+  /**
+   * Returns key of the leaf. It's used for indexing.
+   */
+  getKey(): bigint;
+  /**
+   * Serializes the leaf into a buffer.
+   */
+  toBuffer(): Buffer;
+  /**
+   * Returns true if the leaf is empty.
+   */
+  isEmpty(): boolean;
+}
+
+/**
+ * Preimage of an indexed merkle tree leaf.
+ */
+export interface IndexedTreeLeafPreimage {
+  /**
+   * Returns key of the leaf corresponding to this preimage.
+   */
+  getKey(): bigint;
+  /**
+   * Returns the key of the next leaf.
+   */
+  getNextKey(): bigint;
+  /**
+   * Returns the index of the next leaf.
+   */
+  getNextIndex(): bigint;
+
+  /**
+   * Returns the preimage as a leaf.
+   */
+  asLeaf(): IndexedTreeLeaf;
+  /**
+   * Serializes the preimage into a buffer.
+   */
+  toBuffer(): Buffer;
+  /**
+   * Serializes the preimage to an array of buffers for hashing.
+   */
+  toHashInputs(): Buffer[];
+}

yarn-project/merkle-tree/package.json

@@ -40,6 +40,7 @@
     "tslib": "^2.4.0"
   },
   "devDependencies": {
+    "@aztec/circuits.js": "workspace:^",
     "@jest/globals": "^29.5.0",
     "@types/jest": "^29.5.0",
     "@types/levelup": "^5.1.2",

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

@@ -4,7 +4,7 @@ export * from './interfaces/merkle_tree.js';
 export * from './interfaces/update_only_tree.js';
 export * from './pedersen.js';
 export * from './sparse_tree/sparse_tree.js';
-export { LowLeafWitnessData, StandardIndexedTree } from './standard_indexed_tree/standard_indexed_tree.js';
+export { StandardIndexedTree } from './standard_indexed_tree/standard_indexed_tree.js';
 export * from './standard_tree/standard_tree.js';
 export { INITIAL_LEAF } from './tree_base.js';
 export { newTree } from './new_tree.js';

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

@@ -1,8 +1,26 @@
-import { LeafData, SiblingPath } from '@aztec/types';
+import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { SiblingPath } from '@aztec/types';
 
-import { LowLeafWitnessData } from '../index.js';
 import { AppendOnlyTree } from './append_only_tree.js';
 
+/**
+ * All of the data to be return during batch insertion.
+ */
+export interface LowLeafWitnessData<N extends number> {
+  /**
+   * Preimage of the low nullifier that proves non membership.
+   */
+  leafPreimage: IndexedTreeLeafPreimage;
+  /**
+   * Sibling path to prove membership of low nullifier.
+   */
+  siblingPath: SiblingPath<N>;
+  /**
+   * The index of low nullifier.
+   */
+  index: bigint;
+}
+
 /**
  * The result of a batch insertion in an indexed merkle tree.
  */
@@ -35,27 +53,30 @@ export interface IndexedTree extends AppendOnlyTree {
    * @param includeUncommitted - If true, the uncommitted changes are included in the search.
    * @returns The found leaf index and a flag indicating if the corresponding leaf's value is equal to `newValue`.
    */
-  findIndexOfPreviousValue(
+  findIndexOfPreviousKey(
     newValue: bigint,
     includeUncommitted: boolean,
-  ): {
-    /**
-     * The index of the found leaf.
-     */
-    index: number;
-    /**
-     * A flag indicating if the corresponding leaf's value is equal to `newValue`.
-     */
-    alreadyPresent: boolean;
-  };
+  ): Promise<
+    | {
+        /**
+         * The index of the found leaf.
+         */
+        index: bigint;
+        /**
+         * A flag indicating if the corresponding leaf's value is equal to `newValue`.
+         */
+        alreadyPresent: boolean;
+      }
+    | undefined
+  >;
 
   /**
-   * Gets the latest LeafData copy.
-   * @param index - Index of the leaf of which to obtain the LeafData copy.
+   * Gets the latest LeafPreimage copy.
+   * @param index - Index of the leaf of which to obtain the LeafPreimage copy.
    * @param includeUncommitted - If true, the uncommitted changes are included in the search.
-   * @returns A copy of the leaf data at the given index or undefined if the leaf was not found.
+   * @returns A copy of the leaf preimage at the given index or undefined if the leaf was not found.
    */
-  getLatestLeafDataCopy(index: number, includeUncommitted: boolean): LeafData | undefined;
+  getLatestLeafPreimageCopy(index: bigint, includeUncommitted: boolean): Promise<IndexedTreeLeafPreimage | undefined>;
 
   /**
    * Batch insert multiple leaves into the tree.

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

@@ -49,4 +49,12 @@ export interface MerkleTree extends SiblingPathSource {
    * @param includeUncommitted - Set to true to include uncommitted updates in the data set.
    */
   getLeafValue(index: bigint, includeUncommitted: boolean): Promise<Buffer | 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): Promise<bigint | undefined>;
 }

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

@@ -1,5 +1,3 @@
-import { LeafData } from '@aztec/types';
-
 import { TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js';
 import { MerkleTree } from './merkle_tree.js';
 
@@ -12,6 +10,5 @@ export interface UpdateOnlyTree extends MerkleTree, TreeSnapshotBuilder {
    * @param leaf - The leaf value to be updated.
    * @param index - The leaf to be updated.
    */
-  // TODO: Make this strictly a Buffer
-  updateLeaf(leaf: Buffer | LeafData, index: bigint): Promise<void>;
+  updateLeaf(leaf: Buffer, index: bigint): Promise<void>;
 }