replaced MerkleTree.getTreeRoots with getStateReference

yarn-project/sequencer-client/src/sequencer/utils.ts

@@ -1,11 +1,4 @@
-import {
-  AppendOnlyTreeSnapshot,
-  Fr,
-  GlobalVariables,
-  Header,
-  PartialStateReference,
-  StateReference,
-} from '@aztec/circuits.js';
+import { AppendOnlyTreeSnapshot, GlobalVariables, Header } from '@aztec/circuits.js';
 import { MerkleTreeOperations } from '@aztec/world-state';
 
 /**
@@ -15,19 +8,11 @@ export async function buildInitialHeader(
   db: MerkleTreeOperations,
   prevBlockGlobalVariables: GlobalVariables = GlobalVariables.empty(), // TODO(benesjan): this should most likely be removed
 ) {
-  const roots = await db.getTreeRoots();
+  const state = await db.getStateReference();
   return new Header(
     AppendOnlyTreeSnapshot.empty(), // TODO(benesjan): is it correct that last archive is 0?
     Buffer.alloc(32, 0),
-    new StateReference(
-      new AppendOnlyTreeSnapshot(Fr.fromBuffer(roots.l1Tol2MessageTreeRoot), 0),
-      new PartialStateReference(
-        new AppendOnlyTreeSnapshot(Fr.fromBuffer(roots.noteHashTreeRoot), 0),
-        new AppendOnlyTreeSnapshot(Fr.fromBuffer(roots.nullifierTreeRoot), 0),
-        new AppendOnlyTreeSnapshot(Fr.fromBuffer(roots.contractDataTreeRoot), 0),
-        new AppendOnlyTreeSnapshot(Fr.fromBuffer(roots.publicDataTreeRoot), 0),
-      ),
-    ),
+    state,
     prevBlockGlobalVariables,
   );
 }

yarn-project/world-state/src/world-state-db/merkle_tree_operations.ts

@@ -1,5 +1,5 @@
 import { L2Block, MerkleTreeId } from '@aztec/circuit-types';
-import { NullifierLeafPreimage } from '@aztec/circuits.js';
+import { NullifierLeafPreimage, StateReference } from '@aztec/circuits.js';
 import { Fr } from '@aztec/foundation/fields';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
@@ -34,25 +34,6 @@ export interface TreeInfo {
   depth: number;
 }
 
-/**
- * The current roots of the commitment trees
- */
-// TODO(benesjan): try getting rid of this
-export type CurrentTreeRoots = {
-  /** Note Hash Tree root. */
-  noteHashTreeRoot: Buffer;
-  /** Contract data tree root. */
-  contractDataTreeRoot: Buffer;
-  /** L1 to L2 Messages data tree root. */
-  l1Tol2MessageTreeRoot: Buffer;
-  /** Nullifier data tree root. */
-  nullifierTreeRoot: Buffer;
-  /** Archive root. */
-  archiveRoot: Buffer;
-  /** Public data tree root */
-  publicDataTreeRoot: Buffer;
-};
-
 /**
  * Defines the interface for operations on a set of Merkle Trees.
  */
@@ -71,9 +52,9 @@ export interface MerkleTreeOperations {
   getTreeInfo(treeId: MerkleTreeId): Promise<TreeInfo>;
 
   /**
-   * Gets the current roots of the commitment trees.
+   * Gets the current state reference.
    */
-  getTreeRoots(): Promise<CurrentTreeRoots>;
+  getStateReference(): Promise<StateReference>;
 
   /**
    * Gets sibling path for a leaf.

yarn-project/world-state/src/world-state-db/merkle_tree_operations_facade.ts

@@ -1,12 +1,12 @@
 import { L2Block, MerkleTreeId } from '@aztec/circuit-types';
-import { NullifierLeafPreimage } from '@aztec/circuits.js';
+import { NullifierLeafPreimage, StateReference } from '@aztec/circuits.js';
 import { Fr } from '@aztec/foundation/fields';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 import { BatchInsertionResult } from '@aztec/merkle-tree';
 import { SiblingPath } from '@aztec/types/membership';
 
 import { MerkleTreeDb } from './merkle_tree_db.js';
-import { CurrentTreeRoots, HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
+import { HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
 
 /**
  * Wraps a MerkleTreeDbOperations to call all functions with a preset includeUncommitted flag.
@@ -25,11 +25,11 @@ export class MerkleTreeOperationsFacade implements MerkleTreeOperations {
   }
 
   /**
-   * Get the current roots of the commitment trees.
-   * @returns The current roots of the trees.
+   * Get the current state reference.
+   * @returns The current state reference.
    */
-  getTreeRoots(): Promise<CurrentTreeRoots> {
-    return this.trees.getTreeRoots(this.includeUncommitted);
+  getStateReference(): Promise<StateReference> {
+    return this.trees.getStateReference(this.includeUncommitted);
   }
 
   /**

yarn-project/world-state/src/world-state-db/merkle_tree_snapshot_operations_facade.ts

@@ -1,11 +1,11 @@
 import { MerkleTreeId } from '@aztec/circuit-types';
-import { Fr } from '@aztec/circuits.js';
+import { AppendOnlyTreeSnapshot, Fr, PartialStateReference, StateReference } from '@aztec/circuits.js';
 import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 import { BatchInsertionResult, IndexedTreeSnapshot, TreeSnapshot } from '@aztec/merkle-tree';
 import { SiblingPath } from '@aztec/types/membership';
 
 import { MerkleTreeDb } from './merkle_tree_db.js';
-import { CurrentTreeRoots, HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
+import { HandleL2BlockResult, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
 
 /**
  * Merkle tree operations on readonly tree snapshots.
@@ -86,7 +86,7 @@ export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations
     };
   }
 
-  async getTreeRoots(): Promise<CurrentTreeRoots> {
+  async getStateReference(): Promise<StateReference> {
     const snapshots = await Promise.all([
       this.#getTreeSnapshot(MerkleTreeId.CONTRACT_TREE),
       this.#getTreeSnapshot(MerkleTreeId.NULLIFIER_TREE),
@@ -96,14 +96,30 @@ export class MerkleTreeSnapshotOperationsFacade implements MerkleTreeOperations
       this.#getTreeSnapshot(MerkleTreeId.ARCHIVE),
     ]);
 
-    return {
-      archiveRoot: snapshots[MerkleTreeId.ARCHIVE].getRoot(),
-      contractDataTreeRoot: snapshots[MerkleTreeId.CONTRACT_TREE].getRoot(),
-      l1Tol2MessageTreeRoot: snapshots[MerkleTreeId.L1_TO_L2_MESSAGE_TREE].getRoot(),
-      noteHashTreeRoot: snapshots[MerkleTreeId.NOTE_HASH_TREE].getRoot(),
-      nullifierTreeRoot: snapshots[MerkleTreeId.NULLIFIER_TREE].getRoot(),
-      publicDataTreeRoot: snapshots[MerkleTreeId.PUBLIC_DATA_TREE].getRoot(),
-    };
+    return new StateReference(
+      new AppendOnlyTreeSnapshot(
+        Fr.fromBuffer(snapshots[MerkleTreeId.L1_TO_L2_MESSAGE_TREE].getRoot()),
+        Number(snapshots[MerkleTreeId.L1_TO_L2_MESSAGE_TREE].getNumLeaves()),
+      ),
+      new PartialStateReference(
+        new AppendOnlyTreeSnapshot(
+          Fr.fromBuffer(snapshots[MerkleTreeId.NOTE_HASH_TREE].getRoot()),
+          Number(snapshots[MerkleTreeId.NOTE_HASH_TREE].getNumLeaves()),
+        ),
+        new AppendOnlyTreeSnapshot(
+          Fr.fromBuffer(snapshots[MerkleTreeId.NULLIFIER_TREE].getRoot()),
+          Number(snapshots[MerkleTreeId.NULLIFIER_TREE].getNumLeaves()),
+        ),
+        new AppendOnlyTreeSnapshot(
+          Fr.fromBuffer(snapshots[MerkleTreeId.CONTRACT_TREE].getRoot()),
+          Number(snapshots[MerkleTreeId.CONTRACT_TREE].getNumLeaves()),
+        ),
+        new AppendOnlyTreeSnapshot(
+          Fr.fromBuffer(snapshots[MerkleTreeId.PUBLIC_DATA_TREE].getRoot()),
+          Number(snapshots[MerkleTreeId.PUBLIC_DATA_TREE].getNumLeaves()),
+        ),
+      ),
+    );
   }
 
   appendLeaves(): Promise<void> {

yarn-project/world-state/src/world-state-db/merkle_trees.ts

@@ -1,6 +1,7 @@
 import { L2Block, MerkleTreeId } from '@aztec/circuit-types';
 import {
   ARCHIVE_HEIGHT,
+  AppendOnlyTreeSnapshot,
   CONTRACT_TREE_HEIGHT,
   Fr,
   GlobalVariables,
@@ -13,8 +14,10 @@ import {
   NullifierLeafPreimage,
   PUBLIC_DATA_SUBTREE_HEIGHT,
   PUBLIC_DATA_TREE_HEIGHT,
+  PartialStateReference,
   PublicDataTreeLeaf,
   PublicDataTreeLeafPreimage,
+  StateReference,
 } from '@aztec/circuits.js';
 import { computeBlockHash, computeGlobalsHash } from '@aztec/circuits.js/abis';
 import { Committable } from '@aztec/foundation/committable';
@@ -37,13 +40,7 @@ import { Hasher } from '@aztec/types/interfaces';
 import { SiblingPath } from '@aztec/types/membership';
 
 import { INITIAL_NULLIFIER_TREE_SIZE, INITIAL_PUBLIC_DATA_TREE_SIZE, MerkleTreeDb } from './merkle_tree_db.js';
-import {
-  CurrentTreeRoots,
-  HandleL2BlockResult,
-  IndexedTreeId,
-  MerkleTreeOperations,
-  TreeInfo,
-} from './merkle_tree_operations.js';
+import { HandleL2BlockResult, IndexedTreeId, MerkleTreeOperations, TreeInfo } from './merkle_tree_operations.js';
 import { MerkleTreeOperationsFacade } from './merkle_tree_operations_facade.js';
 
 /**
@@ -234,28 +231,33 @@ export class MerkleTrees implements MerkleTreeDb {
   }
 
   /**
-   * Get the current roots of the commitment trees.
+   * Get the current state reference
    * @param includeUncommitted - Indicates whether to include uncommitted data.
-   * @returns The current roots of the trees.
+   * @returns The current state reference
    */
-  public async getTreeRoots(includeUncommitted: boolean): Promise<CurrentTreeRoots> {
-    const roots = await this.synchronize(() => Promise.resolve(this._getAllTreeRoots(includeUncommitted)));
-
-    return {
-      noteHashTreeRoot: roots[0],
-      nullifierTreeRoot: roots[1],
-      contractDataTreeRoot: roots[2],
-      l1Tol2MessageTreeRoot: roots[3],
-      publicDataTreeRoot: roots[4],
-      archiveRoot: roots[5],
-    };
+  public async getStateReference(includeUncommitted: boolean): Promise<StateReference> {
+    const roots = (await this.synchronize(() => Promise.resolve(this._getAllTreeRoots(includeUncommitted)))).map(root =>
+      Fr.fromBuffer(root),
+    );
+    const numLeaves = await this.synchronize(() => Promise.resolve(this._getAllTreeNumLeaves(includeUncommitted)));
+
+    return new StateReference(
+      new AppendOnlyTreeSnapshot(roots[3], numLeaves[3]),
+      new PartialStateReference(
+        new AppendOnlyTreeSnapshot(roots[0], numLeaves[0]),
+        new AppendOnlyTreeSnapshot(roots[1], numLeaves[1]),
+        new AppendOnlyTreeSnapshot(roots[2], numLeaves[2]),
+        new AppendOnlyTreeSnapshot(roots[4], numLeaves[4]),
+      ),
+    );
   }
 
   private async _getCurrentBlockHash(globalsHash: Fr, includeUncommitted: boolean): Promise<Fr> {
     const roots = (await this._getAllTreeRoots(includeUncommitted)).map(root => Fr.fromBuffer(root));
     return computeBlockHash(globalsHash, roots[0], roots[1], roots[2], roots[3], roots[4]);
   }
 
+  // TODO(benesjan): This is ugly. Try nuking it.
   private _getAllTreeRoots(includeUncommitted: boolean): Promise<Buffer[]> {
     const roots = [
       MerkleTreeId.NOTE_HASH_TREE,
@@ -269,6 +271,20 @@ export class MerkleTrees implements MerkleTreeDb {
     return Promise.resolve(roots);
   }
 
+  // TODO(benesjan): This is ugly. Try nuking it.
+  private _getAllTreeNumLeaves(includeUncommitted: boolean): Promise<number[]> {
+    const sizes = [
+      MerkleTreeId.NOTE_HASH_TREE,
+      MerkleTreeId.NULLIFIER_TREE,
+      MerkleTreeId.CONTRACT_TREE,
+      MerkleTreeId.L1_TO_L2_MESSAGE_TREE,
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      MerkleTreeId.ARCHIVE,
+    ].map(tree => Number(this.trees[tree].getNumLeaves(includeUncommitted)));
+
+    return Promise.resolve(sizes);
+  }
+
   /**
    * Gets the value at the given index.
    * @param treeId - The ID of the tree to get the leaf value from.