@aztec/world-state 0.30.1 → 0.32.0

package/src/world-state-db/merkle_tree_operations.ts

@@ -1,8 +1,8 @@
 import { L2Block, MerkleTreeId, SiblingPath } from '@aztec/circuit-types';
-import { Header, NullifierLeafPreimage, StateReference } from '@aztec/circuits.js';
+import { Fr, Header, NullifierLeafPreimage, StateReference } from '@aztec/circuits.js';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
-import { BatchInsertionResult } from '@aztec/merkle-tree';
+import { AppendOnlyTree, BatchInsertionResult, IndexedTree } from '@aztec/merkle-tree';
 
 /**
  * Type alias for the nullifier tree ID.
@@ -32,6 +32,24 @@ export interface TreeInfo {
   depth: number;
 }
 
+export type MerkleTreeMap = {
+  [MerkleTreeId.NULLIFIER_TREE]: IndexedTree;
+  [MerkleTreeId.NOTE_HASH_TREE]: AppendOnlyTree<Fr>;
+  [MerkleTreeId.PUBLIC_DATA_TREE]: IndexedTree;
+  [MerkleTreeId.L1_TO_L2_MESSAGE_TREE]: AppendOnlyTree<Fr>;
+  [MerkleTreeId.ARCHIVE]: AppendOnlyTree<Fr>;
+};
+
+type LeafTypes = {
+  [MerkleTreeId.NULLIFIER_TREE]: Buffer;
+  [MerkleTreeId.NOTE_HASH_TREE]: Fr;
+  [MerkleTreeId.PUBLIC_DATA_TREE]: Buffer;
+  [MerkleTreeId.L1_TO_L2_MESSAGE_TREE]: Fr;
+  [MerkleTreeId.ARCHIVE]: Fr;
+};
+
+export type MerkleTreeLeafType<ID extends MerkleTreeId> = LeafTypes[ID];
+
 /**
  * Defines the interface for operations on a set of Merkle Trees.
  */
@@ -41,7 +59,7 @@ export interface MerkleTreeOperations {
    * @param treeId - The tree to be updated.
    * @param leaves - The set of leaves to be appended.
    */
-  appendLeaves(treeId: MerkleTreeId, leaves: Buffer[]): Promise<void>;
+  appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void>;
 
   /**
    * Returns information about the given tree.
@@ -71,8 +89,8 @@ export interface MerkleTreeOperations {
    * @param treeId - The tree for which the previous value index is required.
    * @param value - The value to be queried.
    */
-  getPreviousValueIndex(
-    treeId: IndexedTreeId,
+  getPreviousValueIndex<ID extends IndexedTreeId>(
+    treeId: ID,
     value: bigint,
   ): Promise<
     | {
@@ -93,7 +111,7 @@ export interface MerkleTreeOperations {
    * @param treeId - The tree for which leaf data should be returned.
    * @param index - The index of the leaf required.
    */
-  getLeafPreimage(treeId: IndexedTreeId, index: bigint): Promise<IndexedTreeLeafPreimage | undefined>;
+  getLeafPreimage<ID extends IndexedTreeId>(treeId: ID, index: bigint): Promise<IndexedTreeLeafPreimage | undefined>;
 
   /**
    * Update the leaf data at the given index.
@@ -101,21 +119,36 @@ export interface MerkleTreeOperations {
    * @param leaf - The updated leaf value.
    * @param index - The index of the leaf to be updated.
    */
-  updateLeaf(treeId: IndexedTreeId, leaf: NullifierLeafPreimage | Buffer, index: bigint): Promise<void>;
+  updateLeaf<ID extends IndexedTreeId>(treeId: ID, leaf: NullifierLeafPreimage | Buffer, index: bigint): Promise<void>;
 
   /**
    * Returns the index containing a leaf value.
    * @param treeId - The tree for which the index should be returned.
    * @param value - The value to search for in the tree.
    */
-  findLeafIndex(treeId: MerkleTreeId, value: Buffer): Promise<bigint | undefined>;
+  findLeafIndex<ID extends MerkleTreeId>(treeId: ID, value: MerkleTreeLeafType<ID>): Promise<bigint | undefined>;
+
+  /**
+   * Returns the first index containing a leaf value after `startIndex`.
+   * @param treeId - The tree for which the index should be returned.
+   * @param value - The value to search for in the tree.
+   * @param startIndex - The index to start searching from (used when skipping nullified messages)
+   */
+  findLeafIndexAfter<ID extends MerkleTreeId>(
+    treeId: ID,
+    value: MerkleTreeLeafType<ID>,
+    startIndex: bigint,
+  ): Promise<bigint | undefined>;
 
   /**
    * Gets the value for a leaf in the tree.
    * @param treeId - The tree for which the index should be returned.
    * @param index - The index of the leaf.
    */
-  getLeafValue(treeId: MerkleTreeId, index: bigint): Promise<Buffer | undefined>;
+  getLeafValue<ID extends MerkleTreeId>(
+    treeId: ID,
+    index: bigint,
+  ): Promise<MerkleTreeLeafType<typeof treeId> | undefined>;
 
   /**
    * Inserts the block hash into the archive.
@@ -131,8 +164,8 @@ export interface MerkleTreeOperations {
    * @param subtreeHeight - Height of the subtree.
    * @returns The witness data for the leaves to be updated when inserting the new ones.
    */
-  batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number>(
-    treeId: MerkleTreeId,
+  batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number, ID extends IndexedTreeId>(
+    treeId: ID,
     leaves: Buffer[],
     subtreeHeight: number,
   ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>>;
@@ -140,8 +173,9 @@ export interface MerkleTreeOperations {
   /**
    * Handles a single L2 block (i.e. Inserts the new note hashes into the merkle tree).
    * @param block - The L2 block to handle.
+   * @param l1ToL2Messages - The L1 to L2 messages for the block.
    */
-  handleL2Block(block: L2Block): Promise<HandleL2BlockResult>;
+  handleL2BlockAndMessages(block: L2Block, l1ToL2Messages: Fr[]): Promise<HandleL2BlockAndMessagesResult>;
 
   /**
    * Commits pending changes to the underlying store.
@@ -154,8 +188,8 @@ export interface MerkleTreeOperations {
   rollback(): Promise<void>;
 }
 
-/** Return type for handleL2Block */
-export type HandleL2BlockResult = {
+/** Return type for handleL2BlockAndMessages */
+export type HandleL2BlockAndMessagesResult = {
   /** Whether the block processed was emitted by our sequencer */ isBlockOurs: boolean;
 };
 

package/src/world-state-db/merkle_tree_operations_facade.ts

@@ -1,10 +1,16 @@
 import { L2Block, MerkleTreeId, SiblingPath } from '@aztec/circuit-types';
-import { Header, NullifierLeafPreimage, StateReference } from '@aztec/circuits.js';
+import { Fr, Header, NullifierLeafPreimage, StateReference } from '@aztec/circuits.js';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 import { BatchInsertionResult } from '@aztec/merkle-tree';
 
 import { MerkleTreeDb } from './merkle_tree_db.js';
-import { HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
+import {
+  HandleL2BlockAndMessagesResult,
+  IndexedTreeId,
+  MerkleTreeLeafType,
+  MerkleTreeOperations,
+  TreeInfo,
+} from './merkle_tree_operations.js';
 
 /**
  * Wraps a MerkleTreeDbOperations to call all functions with a preset includeUncommitted flag.
@@ -44,7 +50,7 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @param leaves - The set of leaves to be appended.
    * @returns The tree info of the specified tree.
    */
-  appendLeaves(treeId: MerkleTreeId, leaves: Buffer[]): Promise<void> {
+  appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void> {
     return this.trees.appendLeaves(treeId, leaves);
   }
 
@@ -66,8 +72,8 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @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`.
    */
-  getPreviousValueIndex(
-    treeId: MerkleTreeId.NULLIFIER_TREE,
+  getPreviousValueIndex<ID extends IndexedTreeId>(
+    treeId: ID,
     value: bigint,
   ): Promise<
     | {
@@ -92,7 +98,7 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @param index - The index to insert into.
    * @returns Empty promise.
    */
-  updateLeaf(treeId: MerkleTreeId.NULLIFIER_TREE, leaf: NullifierLeafPreimage, index: bigint): Promise<void> {
+  updateLeaf<ID extends IndexedTreeId>(treeId: ID, leaf: NullifierLeafPreimage, index: bigint): Promise<void> {
     return this.trees.updateLeaf(treeId, leaf, index);
   }
 
@@ -102,8 +108,8 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @param index - The index of the leaf to get.
    * @returns Leaf preimage.
    */
-  async getLeafPreimage(
-    treeId: MerkleTreeId.NULLIFIER_TREE,
+  async getLeafPreimage<ID extends IndexedTreeId>(
+    treeId: ID,
     index: bigint,
   ): Promise<IndexedTreeLeafPreimage | undefined> {
     const preimage = await this.trees.getLeafPreimage(treeId, index, this.includeUncommitted);
@@ -116,10 +122,24 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @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(treeId: MerkleTreeId, value: Buffer): Promise<bigint | undefined> {
+  findLeafIndex<ID extends MerkleTreeId>(treeId: ID, value: MerkleTreeLeafType<ID>): Promise<bigint | undefined> {
     return this.trees.findLeafIndex(treeId, value, this.includeUncommitted);
   }
 
+  /**
+   * Returns the first index containing a leaf value after `startIndex`.
+   * @param treeId - The tree for which the index should be returned.
+   * @param value - The value to search for in the tree.
+   * @param startIndex - The index to start searching from (used when skipping nullified messages)
+   */
+  findLeafIndexAfter<ID extends MerkleTreeId>(
+    treeId: ID,
+    value: MerkleTreeLeafType<ID>,
+    startIndex: bigint,
+  ): Promise<bigint | undefined> {
+    return this.trees.findLeafIndexAfter(treeId, value, startIndex, this.includeUncommitted);
+  }
+
   /**
    * Gets the value at the given index.
    * @param treeId - The ID of the tree to get the leaf value from.
@@ -127,8 +147,13 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @param includeUncommitted - Indicates whether to include uncommitted changes.
    * @returns Leaf value at the given index (undefined if not found).
    */
-  getLeafValue(treeId: MerkleTreeId, index: bigint): Promise<Buffer | undefined> {
-    return this.trees.getLeafValue(treeId, index, this.includeUncommitted);
+  getLeafValue<ID extends MerkleTreeId>(
+    treeId: ID,
+    index: bigint,
+  ): Promise<MerkleTreeLeafType<typeof treeId> | undefined> {
+    return this.trees.getLeafValue(treeId, index, this.includeUncommitted) as Promise<
+      MerkleTreeLeafType<typeof treeId> | undefined
+    >;
   }
 
   /**
@@ -143,10 +168,11 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
   /**
    * Handles a single L2 block (i.e. Inserts the new note hashes into the merkle tree).
    * @param block - The L2 block to handle.
+   * @param l1ToL2Messages - The L1 to L2 messages for the block.
    * @returns Whether the block handled was produced by this same node.
    */
-  public handleL2Block(block: L2Block): Promise<HandleL2BlockResult> {
-    return this.trees.handleL2Block(block);
+  public handleL2BlockAndMessages(block: L2Block, l1ToL2Messages: Fr[]): Promise<HandleL2BlockAndMessagesResult> {
+    return this.trees.handleL2BlockAndMessages(block, l1ToL2Messages);
   }
 
   /**
@@ -173,7 +199,7 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
    * @returns The data for the leaves to be updated when inserting the new ones.
    */
   public batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number>(
-    treeId: MerkleTreeId,
+    treeId: IndexedTreeId,
     leaves: Buffer[],
     subtreeHeight: number,
   ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>> {

package/src/world-state-db/merkle_tree_snapshot_operations_facade.ts

@@ -1,10 +1,16 @@
 import { MerkleTreeId, SiblingPath } from '@aztec/circuit-types';
 import { AppendOnlyTreeSnapshot, Fr, Header, PartialStateReference, StateReference } from '@aztec/circuits.js';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
-import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree';
+import { BatchInsertionResult, IndexedTreeSnapshot } from '@aztec/merkle-tree';
 
-import { MerkleTreeDb } from './merkle_tree_db.js';
-import { HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
+import { MerkleTreeDb, TreeSnapshots } from './merkle_tree_db.js';
+import {
+  HandleL2BlockAndMessagesResult,
+  IndexedTreeId,
+  MerkleTreeLeafType,
+  MerkleTreeOperations,
+  TreeInfo,
+} from './merkle_tree_operations.js';
 
 /**
  * Merkle tree operations on readonly tree snapshots.
@@ -12,42 +18,56 @@ import { HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tr
 export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations {
   #treesDb: MerkleTreeDb;
   #blockNumber: number;
-  #treeSnapshots: ReadonlyArray<TreeSnapshot | IndexedTreeSnapshot> = [];
+  #treeSnapshots: TreeSnapshots = {} as any;
 
   constructor(trees: MerkleTreeDb, blockNumber: number) {
     this.#treesDb = trees;
     this.#blockNumber = blockNumber;
   }
 
-  async #getTreeSnapshot(merkleTreeId: number): Promise<TreeSnapshot | IndexedTreeSnapshot> {
-    if (this.#treeSnapshots[merkleTreeId]) {
-      return this.#treeSnapshots[merkleTreeId];
+  async #getTreeSnapshot(treeId: MerkleTreeId): Promise<TreeSnapshots[typeof treeId]> {
+    if (this.#treeSnapshots[treeId]) {
+      return this.#treeSnapshots[treeId];
     }
 
     this.#treeSnapshots = await this.#treesDb.getSnapshot(this.#blockNumber);
-    return this.#treeSnapshots[merkleTreeId]!;
+    return this.#treeSnapshots[treeId]!;
   }
 
-  async findLeafIndex(treeId: MerkleTreeId, value: Buffer): Promise<bigint | undefined> {
+  async findLeafIndex<ID extends MerkleTreeId>(treeId: ID, value: MerkleTreeLeafType<ID>): Promise<bigint | undefined> {
     const tree = await this.#getTreeSnapshot(treeId);
-    return tree.findLeafIndex(value);
+    // TODO #5448 fix "as any"
+    return tree.findLeafIndex(value as any);
   }
 
-  async getLeafPreimage(
-    treeId: MerkleTreeId.NULLIFIER_TREE,
+  async findLeafIndexAfter<ID extends MerkleTreeId>(
+    treeId: MerkleTreeId,
+    value: MerkleTreeLeafType<ID>,
+    startIndex: bigint,
+  ): Promise<bigint | undefined> {
+    const tree = await this.#getTreeSnapshot(treeId);
+    // TODO #5448 fix "as any"
+    return tree.findLeafIndexAfter(value as any, startIndex);
+  }
+
+  async getLeafPreimage<ID extends IndexedTreeId>(
+    treeId: ID,
     index: bigint,
   ): Promise<IndexedTreeLeafPreimage | undefined> {
     const snapshot = (await this.#getTreeSnapshot(treeId)) as IndexedTreeSnapshot;
     return snapshot.getLatestLeafPreimageCopy(BigInt(index));
   }
 
-  async getLeafValue(treeId: MerkleTreeId, index: bigint): Promise<Buffer | undefined> {
+  async getLeafValue<ID extends MerkleTreeId>(
+    treeId: ID,
+    index: bigint,
+  ): Promise<MerkleTreeLeafType<typeof treeId> | undefined> {
     const snapshot = await this.#getTreeSnapshot(treeId);
-    return snapshot.getLeafValue(BigInt(index));
+    return snapshot.getLeafValue(BigInt(index)) as MerkleTreeLeafType<typeof treeId> | undefined;
   }
 
   async getPreviousValueIndex(
-    treeId: MerkleTreeId.NULLIFIER_TREE,
+    treeId: IndexedTreeId,
     value: bigint,
   ): Promise<
     | {
@@ -130,7 +150,7 @@ export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations
     return Promise.reject(new Error('Tree snapshot operations are read-only'));
   }
 
-  handleL2Block(): Promise<HandleL2BlockResult> {
+  handleL2BlockAndMessages(): Promise<HandleL2BlockAndMessagesResult> {
     return Promise.reject(new Error('Tree snapshot operations are read-only'));
   }
 

package/src/world-state-db/merkle_trees.ts

@@ -11,6 +11,7 @@ import {
   NOTE_HASH_TREE_HEIGHT,
   NULLIFIER_SUBTREE_HEIGHT,
   NULLIFIER_TREE_HEIGHT,
+  NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
   NullifierLeaf,
   NullifierLeafPreimage,
   PUBLIC_DATA_SUBTREE_HEIGHT,
@@ -20,6 +21,7 @@ import {
   PublicDataTreeLeafPreimage,
   StateReference,
 } from '@aztec/circuits.js';
+import { padArrayEnd } from '@aztec/foundation/collection';
 import { SerialQueue } from '@aztec/foundation/fifo';
 import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
@@ -38,15 +40,35 @@ import {
 } from '@aztec/merkle-tree';
 import { Hasher } from '@aztec/types/interfaces';
 
-import { INITIAL_NULLIFIER_TREE_SIZE, INITIAL_PUBLIC_DATA_TREE_SIZE, MerkleTreeDb } from './merkle_tree_db.js';
-import { HandleL2BlockResult, IndexedTreeId, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
+import {
+  INITIAL_NULLIFIER_TREE_SIZE,
+  INITIAL_PUBLIC_DATA_TREE_SIZE,
+  MerkleTreeDb,
+  TreeSnapshots,
+} from './merkle_tree_db.js';
+import {
+  HandleL2BlockAndMessagesResult,
+  IndexedTreeId,
+  MerkleTreeLeafType,
+  MerkleTreeMap,
+  MerkleTreeOperations,
+  TreeInfo,
+} from './merkle_tree_operations.js';
 import { MerkleTreeOperationsFacade } from './merkle_tree_operations_facade.js';
 
 /**
  * The nullifier tree is an indexed tree.
  */
 class NullifierTree extends StandardIndexedTree {
-  constructor(store: AztecKVStore, hasher: Hasher, name: string, depth: number, size: bigint = 0n, root?: Buffer) {
+  constructor(
+    store: AztecKVStore,
+    hasher: Hasher,
+    name: string,
+    depth: number,
+    size: bigint = 0n,
+    _noop: any,
+    root?: Buffer,
+  ) {
     super(store, hasher, name, depth, size, NullifierLeafPreimage, NullifierLeaf, root);
   }
 }
@@ -55,7 +77,15 @@ class NullifierTree extends StandardIndexedTree {
  * The public data tree is an indexed tree.
  */
 class PublicDataTree extends StandardIndexedTree {
-  constructor(store: AztecKVStore, hasher: Hasher, name: string, depth: number, size: bigint = 0n, root?: Buffer) {
+  constructor(
+    store: AztecKVStore,
+    hasher: Hasher,
+    name: string,
+    depth: number,
+    size: bigint = 0n,
+    _noop: any,
+    root?: Buffer,
+  ) {
     super(store, hasher, name, depth, size, PublicDataTreeLeafPreimage, PublicDataTreeLeaf, root);
   }
 }
@@ -64,7 +94,8 @@ class PublicDataTree extends StandardIndexedTree {
  * A convenience class for managing multiple merkle trees.
  */
 export class MerkleTrees implements MerkleTreeDb {
-  private trees: (AppendOnlyTree | UpdateOnlyTree)[] = [];
+  // gets initialized in #init
+  private trees: MerkleTreeMap = null as any;
   private jobQueue = new SerialQueue();
 
   private constructor(private store: AztecKVStore, private log: DebugLogger) {}
@@ -94,14 +125,16 @@ export class MerkleTrees implements MerkleTreeDb {
       this.store,
       hasher,
       `${MerkleTreeId[MerkleTreeId.NULLIFIER_TREE]}`,
+      {},
       NULLIFIER_TREE_HEIGHT,
       INITIAL_NULLIFIER_TREE_SIZE,
     );
-    const noteHashTree: AppendOnlyTree = await initializeTree(
+    const noteHashTree: AppendOnlyTree<Fr> = await initializeTree(
       StandardTree,
       this.store,
       hasher,
       `${MerkleTreeId[MerkleTreeId.NOTE_HASH_TREE]}`,
+      Fr,
       NOTE_HASH_TREE_HEIGHT,
     );
     const publicDataTree = await initializeTree(
@@ -109,21 +142,24 @@ export class MerkleTrees implements MerkleTreeDb {
       this.store,
       hasher,
       `${MerkleTreeId[MerkleTreeId.PUBLIC_DATA_TREE]}`,
+      {},
       PUBLIC_DATA_TREE_HEIGHT,
       INITIAL_PUBLIC_DATA_TREE_SIZE,
     );
-    const l1Tol2MessageTree: AppendOnlyTree = await initializeTree(
+    const l1Tol2MessageTree: AppendOnlyTree<Fr> = await initializeTree(
       StandardTree,
       this.store,
       hasher,
       `${MerkleTreeId[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]}`,
+      Fr,
       L1_TO_L2_MSG_TREE_HEIGHT,
     );
-    const archive: AppendOnlyTree = await initializeTree(
+    const archive: AppendOnlyTree<Fr> = await initializeTree(
       StandardTree,
       this.store,
       hasher,
       `${MerkleTreeId[MerkleTreeId.ARCHIVE]}`,
+      Fr,
       ARCHIVE_HEIGHT,
     );
     this.trees = [nullifierTree, noteHashTree, publicDataTree, l1Tol2MessageTree, archive];
@@ -223,7 +259,7 @@ export class MerkleTrees implements MerkleTreeDb {
     treeId: MerkleTreeId,
     index: bigint,
     includeUncommitted: boolean,
-  ): Promise<Buffer | undefined> {
+  ): Promise<MerkleTreeLeafType<typeof treeId> | undefined> {
     return await this.synchronize(() => Promise.resolve(this.trees[treeId].getLeafValue(index, includeUncommitted)));
   }
 
@@ -248,7 +284,7 @@ export class MerkleTrees implements MerkleTreeDb {
    * @param leaves - The leaves to append.
    * @returns Empty promise.
    */
-  public async appendLeaves(treeId: MerkleTreeId, leaves: Buffer[]): Promise<void> {
+  public async appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void> {
     return await this.synchronize(() => this.#appendLeaves(treeId, leaves));
   }
 
@@ -321,14 +357,35 @@ export class MerkleTrees implements MerkleTreeDb {
    * @param includeUncommitted - Indicates whether to include uncommitted data.
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  public async findLeafIndex(
-    treeId: MerkleTreeId,
-    value: Buffer,
+  public async findLeafIndex<ID extends MerkleTreeId>(
+    treeId: ID,
+    value: MerkleTreeLeafType<ID>,
+    includeUncommitted: boolean,
+  ): Promise<bigint | undefined> {
+    return await this.synchronize(() => {
+      const tree = this.trees[treeId];
+      // TODO #5448 fix "as any"
+      return Promise.resolve(tree.findLeafIndex(value as any, includeUncommitted));
+    });
+  }
+
+  /**
+   * Returns the first index containing a leaf value after `startIndex`.
+   * @param treeId - The tree for which the index should be returned.
+   * @param value - The value to search for in the tree.
+   * @param startIndex - The index to start searching from (used when skipping nullified messages)
+   * @param includeUncommitted - Indicates whether to include uncommitted data.
+   */
+  public async findLeafIndexAfter<ID extends MerkleTreeId>(
+    treeId: ID,
+    value: MerkleTreeLeafType<ID>,
+    startIndex: bigint,
     includeUncommitted: boolean,
   ): Promise<bigint | undefined> {
     return await this.synchronize(() => {
       const tree = this.trees[treeId];
-      return Promise.resolve(tree.findLeafIndex(value, includeUncommitted));
+      // TODO #5448 fix "as any"
+      return Promise.resolve(tree.findLeafIndexAfter(value as any, startIndex, includeUncommitted));
     });
   }
 
@@ -346,10 +403,11 @@ export class MerkleTrees implements MerkleTreeDb {
   /**
    * Handles a single L2 block (i.e. Inserts the new note hashes into the merkle tree).
    * @param block - The L2 block to handle.
+   * @param l1ToL2Messages - The L1 to L2 messages for the block.
    * @returns Whether the block handled was produced by this same node.
    */
-  public async handleL2Block(block: L2Block): Promise<HandleL2BlockResult> {
-    return await this.synchronize(() => this.#handleL2Block(block));
+  public async handleL2BlockAndMessages(block: L2Block, l1ToL2Messages: Fr[]): Promise<HandleL2BlockAndMessagesResult> {
+    return await this.synchronize(() => this.#handleL2BlockAndMessages(block, l1ToL2Messages));
   }
 
   /**
@@ -364,7 +422,7 @@ export class MerkleTrees implements MerkleTreeDb {
     SubtreeHeight extends number,
     SubtreeSiblingPathHeight extends number,
   >(
-    treeId: MerkleTreeId,
+    treeId: IndexedTreeId,
     leaves: Buffer[],
     subtreeHeight: SubtreeHeight,
   ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>> {
@@ -394,7 +452,7 @@ export class MerkleTrees implements MerkleTreeDb {
     }
 
     const blockHash = header.hash();
-    await this.#appendLeaves(MerkleTreeId.ARCHIVE, [blockHash.toBuffer()]);
+    await this.#appendLeaves(MerkleTreeId.ARCHIVE, [blockHash]);
   }
 
   /**
@@ -428,20 +486,21 @@ export class MerkleTrees implements MerkleTreeDb {
    * @param leaves - Leaves to append.
    * @returns Empty promise.
    */
-  async #appendLeaves(treeId: MerkleTreeId, leaves: Buffer[]): Promise<void> {
+  async #appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<typeof treeId>[]): Promise<void> {
     const tree = this.trees[treeId];
     if (!('appendLeaves' in tree)) {
       throw new Error('Tree does not support `appendLeaves` method');
     }
-    return await tree.appendLeaves(leaves);
+    // TODO #5448 fix "as any"
+    return await tree.appendLeaves(leaves as any[]);
   }
 
-  async #updateLeaf(treeId: IndexedTreeId, leaf: Buffer, index: bigint): Promise<void> {
+  async #updateLeaf(treeId: IndexedTreeId, leaf: MerkleTreeLeafType<typeof treeId>, index: bigint): Promise<void> {
     const tree = this.trees[treeId];
     if (!('updateLeaf' in tree)) {
       throw new Error('Tree does not support `updateLeaf` method');
     }
-    return await tree.updateLeaf(leaf, index);
+    return await (tree as UpdateOnlyTree<typeof leaf>).updateLeaf(leaf, index);
   }
 
   /**
@@ -449,7 +508,7 @@ export class MerkleTrees implements MerkleTreeDb {
    * @returns Empty promise.
    */
   async #commit(): Promise<void> {
-    for (const tree of this.trees) {
+    for (const tree of Object.values(this.trees)) {
       await tree.commit();
     }
   }
@@ -459,17 +518,31 @@ export class MerkleTrees implements MerkleTreeDb {
    * @returns Empty promise.
    */
   async #rollback(): Promise<void> {
-    for (const tree of this.trees) {
+    for (const tree of Object.values(this.trees)) {
       await tree.rollback();
     }
   }
 
-  public getSnapshot(blockNumber: number) {
-    return Promise.all(this.trees.map(tree => tree.getSnapshot(blockNumber)));
+  public async getSnapshot(blockNumber: number): Promise<TreeSnapshots> {
+    const snapshots = await Promise.all([
+      this.trees[MerkleTreeId.NULLIFIER_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.NOTE_HASH_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.PUBLIC_DATA_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.L1_TO_L2_MESSAGE_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.ARCHIVE].getSnapshot(blockNumber),
+    ]);
+
+    return {
+      [MerkleTreeId.NULLIFIER_TREE]: snapshots[0],
+      [MerkleTreeId.NOTE_HASH_TREE]: snapshots[1],
+      [MerkleTreeId.PUBLIC_DATA_TREE]: snapshots[2],
+      [MerkleTreeId.L1_TO_L2_MESSAGE_TREE]: snapshots[3],
+      [MerkleTreeId.ARCHIVE]: snapshots[4],
+    };
   }
 
   async #snapshot(blockNumber: number): Promise<void> {
-    for (const tree of this.trees) {
+    for (const tree of Object.values(this.trees)) {
       await tree.snapshot(blockNumber);
     }
   }
@@ -477,8 +550,9 @@ export class MerkleTrees implements MerkleTreeDb {
   /**
    * Handles a single L2 block (i.e. Inserts the new note hashes into the merkle tree).
    * @param l2Block - The L2 block to handle.
+   * @param l1ToL2Messages - The L1 to L2 messages for the block.
    */
-  async #handleL2Block(l2Block: L2Block): Promise<HandleL2BlockResult> {
+  async #handleL2BlockAndMessages(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<HandleL2BlockAndMessagesResult> {
     const treeRootWithIdPairs = [
       [l2Block.header.state.partial.nullifierTree.root, MerkleTreeId.NULLIFIER_TREE],
       [l2Block.header.state.partial.noteHashTree.root, MerkleTreeId.NOTE_HASH_TREE],
@@ -498,15 +572,16 @@ export class MerkleTrees implements MerkleTreeDb {
       this.log(`Block ${l2Block.number} is not ours, rolling back world state and committing state from chain`);
       await this.#rollback();
 
+      // We pad the messages because always a fixed number of messages is inserted and we need
+      // the `nextAvailableLeafIndex` to correctly progress.
+      const l1ToL2MessagesPadded = padArrayEnd(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP);
+
       // Sync the append only trees
       for (const [tree, leaves] of [
         [MerkleTreeId.NOTE_HASH_TREE, l2Block.body.txEffects.flatMap(txEffect => txEffect.noteHashes)],
-        [MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l2Block.body.l1ToL2Messages],
+        [MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l1ToL2MessagesPadded],
       ] as const) {
-        await this.#appendLeaves(
-          tree,
-          leaves.map(fr => fr.toBuffer()),
-        );
+        await this.#appendLeaves(tree, leaves);
       }
 
       // Sync the indexed trees