feat: use tree snapshots in aztec-node/pxe/oracles (#3504)

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

@@ -27,6 +27,7 @@ import {
   ContractDataSource,
   ExtendedContractData,
   GetUnencryptedLogsResponse,
+  INITIAL_L2_BLOCK_NUM,
   L1ToL2MessageAndIndex,
   L1ToL2MessageSource,
   L2Block,
@@ -301,42 +302,59 @@ export class AztecNodeService implements AztecNode {
 
   /**
    * Find the index of the given leaf in the given tree.
+   * @param blockNumber - The block number at which to get the data
    * @param treeId - The tree to search in.
    * @param leafValue - The value to search for
    * @returns The index of the given leaf in the given tree or undefined if not found.
    */
-  public async findLeafIndex(treeId: MerkleTreeId, leafValue: Fr): Promise<bigint | undefined> {
-    const committedDb = await this.#getWorldState();
+  public async findLeafIndex(
+    blockNumber: number | 'latest',
+    treeId: MerkleTreeId,
+    leafValue: Fr,
+  ): Promise<bigint | undefined> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.findLeafIndex(treeId, leafValue.toBuffer());
   }
 
   /**
    * Returns a sibling path for the given index in the contract tree.
+   * @param blockNumber - The block number at which to get the data.
    * @param leafIndex - The index of the leaf for which the sibling path is required.
    * @returns The sibling path for the leaf index.
    */
-  public async getContractSiblingPath(leafIndex: bigint): Promise<SiblingPath<typeof CONTRACT_TREE_HEIGHT>> {
-    const committedDb = await this.#getWorldState();
+  public async getContractSiblingPath(
+    blockNumber: number | 'latest',
+    leafIndex: bigint,
+  ): Promise<SiblingPath<typeof CONTRACT_TREE_HEIGHT>> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.getSiblingPath(MerkleTreeId.CONTRACT_TREE, leafIndex);
   }
 
   /**
    * Returns a sibling path for the given index in the nullifier tree.
+   * @param blockNumber - The block number at which to get the data.
    * @param leafIndex - The index of the leaf for which the sibling path is required.
    * @returns The sibling path for the leaf index.
    */
-  public async getNullifierTreeSiblingPath(leafIndex: bigint): Promise<SiblingPath<typeof NULLIFIER_TREE_HEIGHT>> {
-    const committedDb = await this.#getWorldState();
+  public async getNullifierTreeSiblingPath(
+    blockNumber: number | 'latest',
+    leafIndex: bigint,
+  ): Promise<SiblingPath<typeof NULLIFIER_TREE_HEIGHT>> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, leafIndex);
   }
 
   /**
    * Returns a sibling path for the given index in the data tree.
+   * @param blockNumber - The block number at which to get the data.
    * @param leafIndex - The index of the leaf for which the sibling path is required.
    * @returns The sibling path for the leaf index.
    */
-  public async getNoteHashSiblingPath(leafIndex: bigint): Promise<SiblingPath<typeof NOTE_HASH_TREE_HEIGHT>> {
-    const committedDb = await this.#getWorldState();
+  public async getNoteHashSiblingPath(
+    blockNumber: number | 'latest',
+    leafIndex: bigint,
+  ): Promise<SiblingPath<typeof NOTE_HASH_TREE_HEIGHT>> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.getSiblingPath(MerkleTreeId.NOTE_HASH_TREE, leafIndex);
   }
 
@@ -348,38 +366,50 @@ export class AztecNodeService implements AztecNode {
    */
   public async getL1ToL2MessageAndIndex(messageKey: Fr): Promise<L1ToL2MessageAndIndex> {
     // todo: #697 - make this one lookup.
-    const index = (await this.findLeafIndex(MerkleTreeId.L1_TO_L2_MESSAGES_TREE, messageKey))!;
+    const index = (await this.findLeafIndex('latest', MerkleTreeId.L1_TO_L2_MESSAGES_TREE, messageKey))!;
     const message = await this.l1ToL2MessageSource.getConfirmedL1ToL2Message(messageKey);
     return Promise.resolve(new L1ToL2MessageAndIndex(index, message));
   }
 
   /**
    * Returns a sibling path for a leaf in the committed l1 to l2 data tree.
+   * @param blockNumber - The block number at which to get the data.
    * @param leafIndex - Index of the leaf in the tree.
    * @returns The sibling path.
    */
-  public async getL1ToL2MessageSiblingPath(leafIndex: bigint): Promise<SiblingPath<typeof L1_TO_L2_MSG_TREE_HEIGHT>> {
-    const committedDb = await this.#getWorldState();
+  public async getL1ToL2MessageSiblingPath(
+    blockNumber: number | 'latest',
+    leafIndex: bigint,
+  ): Promise<SiblingPath<typeof L1_TO_L2_MSG_TREE_HEIGHT>> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.getSiblingPath(MerkleTreeId.L1_TO_L2_MESSAGES_TREE, leafIndex);
   }
 
   /**
    * Returns a sibling path for a leaf in the committed blocks tree.
+   * @param blockNumber - The block number at which to get the data.
    * @param leafIndex - Index of the leaf in the tree.
    * @returns The sibling path.
    */
-  public async getBlocksTreeSiblingPath(leafIndex: bigint): Promise<SiblingPath<typeof BLOCKS_TREE_HEIGHT>> {
-    const committedDb = await this.#getWorldState();
+  public async getBlocksTreeSiblingPath(
+    blockNumber: number | 'latest',
+    leafIndex: bigint,
+  ): Promise<SiblingPath<typeof BLOCKS_TREE_HEIGHT>> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.getSiblingPath(MerkleTreeId.BLOCKS_TREE, leafIndex);
   }
 
   /**
    * Returns a sibling path for a leaf in the committed public data tree.
+   * @param blockNumber - The block number at which to get the data.
    * @param leafIndex - Index of the leaf in the tree.
    * @returns The sibling path.
    */
-  public async getPublicDataTreeSiblingPath(leafIndex: bigint): Promise<SiblingPath<typeof PUBLIC_DATA_TREE_HEIGHT>> {
-    const committedDb = await this.#getWorldState();
+  public async getPublicDataTreeSiblingPath(
+    blockNumber: number | 'latest',
+    leafIndex: bigint,
+  ): Promise<SiblingPath<typeof PUBLIC_DATA_TREE_HEIGHT>> {
+    const committedDb = await this.#getWorldState(blockNumber);
     return committedDb.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, leafIndex);
   }
 
@@ -390,17 +420,17 @@ export class AztecNodeService implements AztecNode {
    * @returns The nullifier membership witness (if found).
    */
   public async getNullifierMembershipWitness(
-    blockNumber: number,
+    blockNumber: number | 'latest',
     nullifier: Fr,
   ): Promise<NullifierMembershipWitness | undefined> {
-    const committedDb = await this.#getWorldState();
-    const index = await committedDb.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBuffer());
+    const db = await this.#getWorldState(blockNumber);
+    const index = await db.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBuffer());
     if (!index) {
       return undefined;
     }
 
-    const leafDataPromise = committedDb.getLeafData(MerkleTreeId.NULLIFIER_TREE, Number(index));
-    const siblingPathPromise = committedDb.getSiblingPath<typeof NULLIFIER_TREE_HEIGHT>(
+    const leafDataPromise = db.getLeafData(MerkleTreeId.NULLIFIER_TREE, Number(index));
+    const siblingPathPromise = db.getSiblingPath<typeof NULLIFIER_TREE_HEIGHT>(
       MerkleTreeId.NULLIFIER_TREE,
       BigInt(index),
     );
@@ -429,10 +459,10 @@ export class AztecNodeService implements AztecNode {
    * TODO: This is a confusing behavior and we should eventually address that.
    */
   public async getLowNullifierMembershipWitness(
-    blockNumber: number,
+    blockNumber: number | 'latest',
     nullifier: Fr,
   ): Promise<NullifierMembershipWitness | undefined> {
-    const committedDb = await this.#getWorldState();
+    const committedDb = await this.#getWorldState(blockNumber);
     const { index, alreadyPresent } = await committedDb.getPreviousValueIndex(
       MerkleTreeId.NULLIFIER_TREE,
       nullifier.toBigInt(),
@@ -462,7 +492,7 @@ export class AztecNodeService implements AztecNode {
    * @returns Storage value at the given contract slot (or undefined if not found).
    */
   public async getPublicStorageAt(contract: AztecAddress, slot: Fr): Promise<Fr | undefined> {
-    const committedDb = await this.#getWorldState();
+    const committedDb = await this.#getWorldState('latest');
     const leafIndex = computePublicDataTreeIndex(contract, slot);
     const value = await committedDb.getLeafValue(MerkleTreeId.PUBLIC_DATA_TREE, leafIndex.value);
     return value ? Fr.fromBuffer(value) : undefined;
@@ -473,7 +503,7 @@ export class AztecNodeService implements AztecNode {
    * @returns The current committed roots for the data trees.
    */
   public async getTreeRoots(): Promise<Record<MerkleTreeId, Fr>> {
-    const committedDb = await this.#getWorldState();
+    const committedDb = await this.#getWorldState('latest');
     const getTreeRoot = async (id: MerkleTreeId) => Fr.fromBuffer((await committedDb.getTreeInfo(id)).root);
 
     const [noteHashTree, nullifierTree, contractTree, l1ToL2MessagesTree, blocksTree, publicDataTree] =
@@ -501,7 +531,7 @@ export class AztecNodeService implements AztecNode {
    * @returns The current committed block header.
    */
   public async getBlockHeader(): Promise<BlockHeader> {
-    const committedDb = await this.#getWorldState();
+    const committedDb = await this.#getWorldState('latest');
     const [roots, globalsHash] = await Promise.all([this.getTreeRoots(), committedDb.getLatestGlobalVariablesHash()]);
 
     return new BlockHeader(
@@ -555,24 +585,40 @@ export class AztecNodeService implements AztecNode {
 
   /**
    * Returns an instance of MerkleTreeOperations having first ensured the world state is fully synched
+   * @param blockNumber - The block number at which to get the data.
    * @returns An instance of a committed MerkleTreeOperations
    */
-  async #getWorldState() {
+  async #getWorldState(blockNumber: number | 'latest') {
+    if (typeof blockNumber === 'number' && blockNumber < INITIAL_L2_BLOCK_NUM) {
+      throw new Error('Invalid block number to get world state for: ' + blockNumber);
+    }
+
+    let blockSyncedTo: number = 0;
     try {
       // Attempt to sync the world state if necessary
-      await this.#syncWorldState();
+      blockSyncedTo = await this.#syncWorldState();
     } catch (err) {
       this.log.error(`Error getting world state: ${err}`);
     }
-    return this.worldStateSynchronizer.getCommitted();
+
+    // using a snapshot could be less efficient than using the committed db
+    if (blockNumber === 'latest' || blockNumber === blockSyncedTo) {
+      this.log(`Using committed db for block ${blockNumber}, world state synced upto ${blockSyncedTo}`);
+      return this.worldStateSynchronizer.getCommitted();
+    } else if (blockNumber < blockSyncedTo) {
+      this.log(`Using snapshot for block ${blockNumber}, world state synced upto ${blockSyncedTo}`);
+      return this.worldStateSynchronizer.getSnapshot(blockNumber);
+    } else {
+      throw new Error(`Block ${blockNumber} not yet synced`);
+    }
   }
 
   /**
    * Ensure we fully sync the world state
    * @returns A promise that fulfils once the world state is synced
    */
-  async #syncWorldState() {
+  async #syncWorldState(): Promise<number> {
     const blockSourceHeight = await this.blockSource.getBlockNumber();
-    await this.worldStateSynchronizer.syncImmediate(blockSourceHeight);
+    return this.worldStateSynchronizer.syncImmediate(blockSourceHeight);
   }
 }

yarn-project/pxe/src/contract_tree/index.ts

@@ -146,12 +146,14 @@ export class ContractTree {
    * If the witness hasn't been previously computed, this function will request the contract node
    * to find the contract's index and path in order to create the membership witness.
    *
+   * @param blockNumber - The block number at which to get the data.
+   *
    * @returns A Promise that resolves to the MembershipWitness object for the given contract tree.
    */
-  public async getContractMembershipWitness() {
+  public async getContractMembershipWitness(blockNumber: number | 'latest' = 'latest') {
     const index = await this.getContractIndex();
 
-    const siblingPath = await this.stateInfoProvider.getContractSiblingPath(index);
+    const siblingPath = await this.stateInfoProvider.getContractSiblingPath(blockNumber, index);
     return new MembershipWitness<typeof CONTRACT_TREE_HEIGHT>(
       CONTRACT_TREE_HEIGHT,
       index,
@@ -226,7 +228,7 @@ export class ContractTree {
       const root = await this.getFunctionTreeRoot();
       const newContractData = new NewContractData(completeAddress.address, portalContract, root);
       const commitment = computeContractLeaf(newContractData);
-      this.contractIndex = await this.stateInfoProvider.findLeafIndex(MerkleTreeId.CONTRACT_TREE, commitment);
+      this.contractIndex = await this.stateInfoProvider.findLeafIndex('latest', MerkleTreeId.CONTRACT_TREE, commitment);
       if (this.contractIndex === undefined) {
         throw new Error(
           `Failed to find contract at ${completeAddress.address} with portal ${portalContract} resulting in commitment ${commitment}.`,

yarn-project/pxe/src/kernel_oracle/index.ts

@@ -24,7 +24,7 @@ export class KernelOracle implements ProvingDataOracle {
   }
 
   async getNoteMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>> {
-    const path = await this.node.getNoteHashSiblingPath(leafIndex);
+    const path = await this.node.getNoteHashSiblingPath('latest', leafIndex);
     return new MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>(
       path.pathSize,
       leafIndex,

yarn-project/pxe/src/pxe_service/pxe_service.ts

@@ -237,13 +237,13 @@ export class PXEService implements PXE {
       // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1386)
       // This can always be `uniqueSiloedNoteHash` once notes added from public also include nonces.
       const noteHashToLookUp = nonce.isZero() ? siloedNoteHash : uniqueSiloedNoteHash;
-      const index = await this.node.findLeafIndex(MerkleTreeId.NOTE_HASH_TREE, noteHashToLookUp);
+      const index = await this.node.findLeafIndex('latest', MerkleTreeId.NOTE_HASH_TREE, noteHashToLookUp);
       if (index === undefined) {
         throw new Error('Note does not exist.');
       }
 
       const siloedNullifier = siloNullifier(note.contractAddress, innerNullifier!);
-      const nullifierIndex = await this.node.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, siloedNullifier);
+      const nullifierIndex = await this.node.findLeafIndex('latest', MerkleTreeId.NULLIFIER_TREE, siloedNullifier);
       if (nullifierIndex !== undefined) {
         throw new Error('The note has been destroyed.');
       }

yarn-project/pxe/src/simulator_oracle/index.ts

@@ -115,7 +115,7 @@ export class SimulatorOracle implements DBOracle {
     const messageAndIndex = await this.stateInfoProvider.getL1ToL2MessageAndIndex(msgKey);
     const message = messageAndIndex.message.toFieldArray();
     const index = messageAndIndex.index;
-    const siblingPath = await this.stateInfoProvider.getL1ToL2MessageSiblingPath(index);
+    const siblingPath = await this.stateInfoProvider.getL1ToL2MessageSiblingPath('latest', index);
     return {
       message,
       siblingPath: siblingPath.toFieldArray(),
@@ -129,32 +129,28 @@ export class SimulatorOracle implements DBOracle {
    * @returns - The index of the commitment. Undefined if it does not exist in the tree.
    */
   async getCommitmentIndex(commitment: Fr) {
-    return await this.stateInfoProvider.findLeafIndex(MerkleTreeId.NOTE_HASH_TREE, commitment);
+    return await this.stateInfoProvider.findLeafIndex('latest', MerkleTreeId.NOTE_HASH_TREE, commitment);
   }
 
   async getNullifierIndex(nullifier: Fr) {
-    return await this.stateInfoProvider.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier);
+    return await this.stateInfoProvider.findLeafIndex('latest', MerkleTreeId.NULLIFIER_TREE, nullifier);
   }
 
   public async findLeafIndex(blockNumber: number, treeId: MerkleTreeId, leafValue: Fr): Promise<bigint | undefined> {
-    this.log.warn('Block number ignored in SimulatorOracle.findLeafIndex because archival node is not yet implemented');
-    return await this.stateInfoProvider.findLeafIndex(treeId, leafValue);
+    return await this.stateInfoProvider.findLeafIndex(blockNumber, treeId, leafValue);
   }
 
   public async getSiblingPath(blockNumber: number, treeId: MerkleTreeId, leafIndex: bigint): Promise<Fr[]> {
-    this.log.warn(
-      'Block number ignored in SimulatorOracle.getSiblingPath because archival node is not yet implemented',
-    );
     // @todo Doing a nasty workaround here because of https://github.com/AztecProtocol/aztec-packages/issues/3414
     switch (treeId) {
       case MerkleTreeId.NULLIFIER_TREE:
-        return (await this.stateInfoProvider.getNullifierTreeSiblingPath(leafIndex)).toFieldArray();
+        return (await this.stateInfoProvider.getNullifierTreeSiblingPath(blockNumber, leafIndex)).toFieldArray();
       case MerkleTreeId.NOTE_HASH_TREE:
-        return (await this.stateInfoProvider.getNoteHashSiblingPath(leafIndex)).toFieldArray();
+        return (await this.stateInfoProvider.getNoteHashSiblingPath(blockNumber, leafIndex)).toFieldArray();
       case MerkleTreeId.BLOCKS_TREE:
-        return (await this.stateInfoProvider.getBlocksTreeSiblingPath(leafIndex)).toFieldArray();
+        return (await this.stateInfoProvider.getBlocksTreeSiblingPath(blockNumber, leafIndex)).toFieldArray();
       case MerkleTreeId.PUBLIC_DATA_TREE:
-        return (await this.stateInfoProvider.getPublicDataTreeSiblingPath(leafIndex)).toFieldArray();
+        return (await this.stateInfoProvider.getPublicDataTreeSiblingPath(blockNumber, leafIndex)).toFieldArray();
       default:
         throw new Error('Not implemented');
     }