feat: PXE handles reorgs

yarn-project/archiver/src/archiver/archiver.ts

@@ -2,6 +2,7 @@ import {
   type EncryptedL2Log,
   type FromLogType,
   type GetUnencryptedLogsResponse,
+  type InBlock,
   type InboxLeaf,
   type L1ToL2MessageSource,
   type L2Block,
@@ -589,7 +590,7 @@ export class Archiver implements ArchiveSource {
     }
   }
 
-  public getTxEffect(txHash: TxHash): Promise<TxEffect | undefined> {
+  public getTxEffect(txHash: TxHash) {
     return this.store.getTxEffect(txHash);
   }
 
@@ -933,7 +934,7 @@ class ArchiverStoreHelper
   getBlockHeaders(from: number, limit: number): Promise<Header[]> {
     return this.store.getBlockHeaders(from, limit);
   }
-  getTxEffect(txHash: TxHash): Promise<TxEffect | undefined> {
+  getTxEffect(txHash: TxHash): Promise<InBlock<TxEffect> | undefined> {
     return this.store.getTxEffect(txHash);
   }
   getSettledTxReceipt(txHash: TxHash): Promise<TxReceipt | undefined> {

yarn-project/archiver/src/archiver/archiver_store.ts

@@ -1,6 +1,7 @@
 import {
   type FromLogType,
   type GetUnencryptedLogsResponse,
+  type InBlock,
   type InboxLeaf,
   type L2Block,
   type L2BlockL2Logs,
@@ -79,7 +80,7 @@ export interface ArchiverDataStore {
    * @param txHash - The txHash of the tx corresponding to the tx effect.
    * @returns The requested tx effect (or undefined if not found).
    */
-  getTxEffect(txHash: TxHash): Promise<TxEffect | undefined>;
+  getTxEffect(txHash: TxHash): Promise<InBlock<TxEffect> | undefined>;
 
   /**
    * Gets a receipt of a settled tx.

yarn-project/archiver/src/archiver/archiver_store_test_suite.ts

@@ -1,4 +1,4 @@
-import { InboxLeaf, L2Block, LogId, LogType, TxHash } from '@aztec/circuit-types';
+import { InboxLeaf, L2Block, LogId, LogType, TxHash, wrapInBlock } from '@aztec/circuit-types';
 import '@aztec/circuit-types/jest';
 import {
   AztecAddress,
@@ -191,14 +191,14 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
       });
 
       it.each([
-        () => blocks[0].data.body.txEffects[0],
-        () => blocks[9].data.body.txEffects[3],
-        () => blocks[3].data.body.txEffects[1],
-        () => blocks[5].data.body.txEffects[2],
-        () => blocks[1].data.body.txEffects[0],
+        () => wrapInBlock(blocks[0].data.body.txEffects[0], blocks[0].data),
+        () => wrapInBlock(blocks[9].data.body.txEffects[3], blocks[9].data),
+        () => wrapInBlock(blocks[3].data.body.txEffects[1], blocks[3].data),
+        () => wrapInBlock(blocks[5].data.body.txEffects[2], blocks[5].data),
+        () => wrapInBlock(blocks[1].data.body.txEffects[0], blocks[1].data),
       ])('retrieves a previously stored transaction', async getExpectedTx => {
         const expectedTx = getExpectedTx();
-        const actualTx = await store.getTxEffect(expectedTx.txHash);
+        const actualTx = await store.getTxEffect(expectedTx.data.txHash);
         expect(actualTx).toEqual(expectedTx);
       });
 
@@ -207,16 +207,16 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
       });
 
       it.each([
-        () => blocks[0].data.body.txEffects[0],
-        () => blocks[9].data.body.txEffects[3],
-        () => blocks[3].data.body.txEffects[1],
-        () => blocks[5].data.body.txEffects[2],
-        () => blocks[1].data.body.txEffects[0],
+        () => wrapInBlock(blocks[0].data.body.txEffects[0], blocks[0].data),
+        () => wrapInBlock(blocks[9].data.body.txEffects[3], blocks[9].data),
+        () => wrapInBlock(blocks[3].data.body.txEffects[1], blocks[3].data),
+        () => wrapInBlock(blocks[5].data.body.txEffects[2], blocks[5].data),
+        () => wrapInBlock(blocks[1].data.body.txEffects[0], blocks[1].data),
       ])('tries to retrieves a previously stored transaction after deleted', async getExpectedTx => {
         await store.unwindBlocks(blocks.length, blocks.length);
 
         const expectedTx = getExpectedTx();
-        const actualTx = await store.getTxEffect(expectedTx.txHash);
+        const actualTx = await store.getTxEffect(expectedTx.data.txHash);
         expect(actualTx).toEqual(undefined);
       });
 

yarn-project/archiver/src/archiver/kv_archiver_store/block_store.ts

@@ -1,4 +1,4 @@
-import { Body, L2Block, type TxEffect, type TxHash, TxReceipt } from '@aztec/circuit-types';
+import { Body, type InBlock, L2Block, type TxEffect, type TxHash, TxReceipt } from '@aztec/circuit-types';
 import { AppendOnlyTreeSnapshot, type AztecAddress, Header, INITIAL_L2_BLOCK_NUM } from '@aztec/circuits.js';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { type AztecKVStore, type AztecMap, type AztecSingleton, type Range } from '@aztec/kv-store';
@@ -170,14 +170,22 @@ export class BlockStore {
    * @param txHash - The txHash of the tx corresponding to the tx effect.
    * @returns The requested tx effect (or undefined if not found).
    */
-  getTxEffect(txHash: TxHash): TxEffect | undefined {
+  getTxEffect(txHash: TxHash): InBlock<TxEffect> | undefined {
     const [blockNumber, txIndex] = this.getTxLocation(txHash) ?? [];
     if (typeof blockNumber !== 'number' || typeof txIndex !== 'number') {
       return undefined;
     }
 
     const block = this.getBlock(blockNumber);
-    return block?.data.body.txEffects[txIndex];
+    if (!block) {
+      return undefined;
+    }
+
+    return {
+      data: block.data.body.txEffects[txIndex],
+      l2BlockNumber: block.data.number,
+      l2BlockHash: block.data.hash().toString(),
+    };
   }
 
   /**

yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.ts

@@ -6,7 +6,6 @@ import {
   type L2BlockL2Logs,
   type LogFilter,
   type LogType,
-  type TxEffect,
   type TxHash,
   type TxReceipt,
   type TxScopedL2Log,
@@ -159,7 +158,7 @@ export class KVArchiverDataStore implements ArchiverDataStore {
    * @param txHash - The txHash of the tx corresponding to the tx effect.
    * @returns The requested tx effect (or undefined if not found).
    */
-  getTxEffect(txHash: TxHash): Promise<TxEffect | undefined> {
+  getTxEffect(txHash: TxHash) {
     return Promise.resolve(this.#blockStore.getTxEffect(txHash));
   }
 

yarn-project/archiver/src/archiver/memory_archiver_store/memory_archiver_store.ts

@@ -6,6 +6,7 @@ import {
   ExtendedUnencryptedL2Log,
   type FromLogType,
   type GetUnencryptedLogsResponse,
+  type InBlock,
   type InboxLeaf,
   type L2Block,
   type L2BlockL2Logs,
@@ -17,6 +18,7 @@ import {
   TxReceipt,
   TxScopedL2Log,
   type UnencryptedL2BlockL2Logs,
+  wrapInBlock,
 } from '@aztec/circuit-types';
 import {
   type ContractClassPublic,
@@ -50,7 +52,7 @@ export class MemoryArchiverStore implements ArchiverDataStore {
   /**
    * An array containing all the tx effects in the L2 blocks that have been fetched so far.
    */
-  private txEffects: TxEffect[] = [];
+  private txEffects: InBlock<TxEffect>[] = [];
 
   private noteEncryptedLogsPerBlock: Map<number, EncryptedNoteL2BlockL2Logs> = new Map();
 
@@ -181,7 +183,7 @@ export class MemoryArchiverStore implements ArchiverDataStore {
 
     this.lastL1BlockNewBlocks = blocks[blocks.length - 1].l1.blockNumber;
     this.l2Blocks.push(...blocks);
-    this.txEffects.push(...blocks.flatMap(b => b.data.body.txEffects));
+    this.txEffects.push(...blocks.flatMap(b => b.data.body.txEffects.map(txEffect => wrapInBlock(txEffect, b.data))));
 
     return Promise.resolve(true);
   }
@@ -365,8 +367,8 @@ export class MemoryArchiverStore implements ArchiverDataStore {
    * @param txHash - The txHash of the tx effect.
    * @returns The requested tx effect.
    */
-  public getTxEffect(txHash: TxHash): Promise<TxEffect | undefined> {
-    const txEffect = this.txEffects.find(tx => tx.txHash.equals(txHash));
+  public getTxEffect(txHash: TxHash): Promise<InBlock<TxEffect> | undefined> {
+    const txEffect = this.txEffects.find(tx => tx.data.txHash.equals(txHash));
     return Promise.resolve(txEffect);
   }
 

yarn-project/archiver/src/test/mock_l2_block_source.ts

@@ -119,8 +119,14 @@ export class MockL2BlockSource implements L2BlockSource {
    * @returns The requested tx effect.
    */
   public getTxEffect(txHash: TxHash) {
-    const txEffect = this.l2Blocks.flatMap(b => b.body.txEffects).find(tx => tx.txHash.equals(txHash));
-    return Promise.resolve(txEffect);
+    const match = this.l2Blocks
+      .flatMap(b => b.body.txEffects.map(tx => [tx, b] as const))
+      .find(([tx]) => tx.txHash.equals(txHash));
+    if (!match) {
+      return Promise.resolve(undefined);
+    }
+    const [txEffect, block] = match;
+    return Promise.resolve({ data: txEffect, l2BlockNumber: block.number, l2BlockHash: block.hash().toString() });
   }
 
   /**

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

@@ -6,6 +6,7 @@ import {
   type EpochProofQuote,
   type FromLogType,
   type GetUnencryptedLogsResponse,
+  type InBlock,
   type L1ToL2MessageSource,
   type L2Block,
   type L2BlockL2Logs,
@@ -32,6 +33,7 @@ import {
   type TxValidator,
   type WorldStateSynchronizer,
   tryStop,
+  wrapInBlock,
 } from '@aztec/circuit-types';
 import {
   type ARCHIVE_HEIGHT,
@@ -117,14 +119,18 @@ export class AztecNodeService implements AztecNode {
     this.log.info(message);
   }
 
-  addEpochProofQuote(quote: EpochProofQuote): Promise<void> {
+  public addEpochProofQuote(quote: EpochProofQuote): Promise<void> {
     return Promise.resolve(this.p2pClient.addEpochProofQuote(quote));
   }
 
-  getEpochProofQuotes(epoch: bigint): Promise<EpochProofQuote[]> {
+  public getEpochProofQuotes(epoch: bigint): Promise<EpochProofQuote[]> {
     return this.p2pClient.getEpochProofQuotes(epoch);
   }
 
+  public getL2Tips() {
+    return this.blockSource.getL2Tips();
+  }
+
   /**
    * initializes the Aztec Node, wait for component to sync.
    * @param config - The configuration to be used by the aztec node.
@@ -372,7 +378,7 @@ export class AztecNodeService implements AztecNode {
     return txReceipt;
   }
 
-  public getTxEffect(txHash: TxHash): Promise<TxEffect | undefined> {
+  public getTxEffect(txHash: TxHash): Promise<InBlock<TxEffect> | undefined> {
     return this.blockSource.getTxEffect(txHash);
   }
 
@@ -426,6 +432,33 @@ export class AztecNodeService implements AztecNode {
     return await Promise.all(leafValues.map(leafValue => committedDb.findLeafIndex(treeId, leafValue.toBuffer())));
   }
 
+  public async findLeavesIndexesWithApproxBlockNumber(
+    maxBlockNumber: L2BlockNumber,
+    minBlockNumber: number,
+    treeId: MerkleTreeId,
+    leafValues: Fr[],
+  ): Promise<(InBlock<bigint> | undefined)[]> {
+    const initialBlockNumber = maxBlockNumber === 'latest' ? await this.getBlockNumber() : maxBlockNumber;
+    return await Promise.all(
+      leafValues.map(async leafValue => {
+        let blockNumber = initialBlockNumber;
+        let treeSnapshot = this.worldStateSynchronizer.getSnapshot(initialBlockNumber);
+        let index = await treeSnapshot.findLeafIndex(treeId, leafValue.toBuffer());
+        if (!index) {
+          return undefined;
+        }
+        let meta = await treeSnapshot.getTreeInfo(treeId);
+        while (meta.size > index && blockNumber >= minBlockNumber) {
+          blockNumber -= 1;
+          treeSnapshot = this.worldStateSynchronizer.getSnapshot(blockNumber);
+          meta = await treeSnapshot.getTreeInfo(treeId);
+        }
+        const block = await this.getBlock(blockNumber + 1);
+        return wrapInBlock(index, block!);
+      }),
+    );
+  }
+
   /**
    * Returns a sibling path for the given index in the nullifier tree.
    * @param blockNumber - The block number at which to get the data.
@@ -708,7 +741,7 @@ export class AztecNodeService implements AztecNode {
    * Returns the currently committed block header, or the initial header if no blocks have been produced.
    * @returns The current committed block header.
    */
-  public async getHeader(blockNumber: L2BlockNumber = 'latest'): Promise<Header> {
+  public async getBlockHeader(blockNumber: L2BlockNumber = 'latest'): Promise<Header> {
     return (
       (await this.getBlock(blockNumber === 'latest' ? -1 : blockNumber))?.header ??
       this.worldStateSynchronizer.getCommitted().getInitialHeader()

yarn-project/aztec.js/src/contract/sent_tx.ts

@@ -78,7 +78,7 @@ export class SentTx {
     }
     if (opts?.debug) {
       const txHash = await this.getTxHash();
-      const tx = (await this.pxe.getTxEffect(txHash))!;
+      const { data: tx } = (await this.pxe.getTxEffect(txHash))!;
       receipt.debugInfo = {
         noteHashes: tx.noteHashes,
         nullifiers: tx.nullifiers,

yarn-project/aztec.js/src/wallet/base_wallet.ts

@@ -13,7 +13,6 @@ import {
   type SiblingPath,
   type SyncStatus,
   type Tx,
-  type TxEffect,
   type TxExecutionRequest,
   type TxHash,
   type TxProvingResult,
@@ -122,7 +121,7 @@ export abstract class BaseWallet implements Wallet {
   sendTx(tx: Tx): Promise<TxHash> {
     return this.pxe.sendTx(tx);
   }
-  getTxEffect(txHash: TxHash): Promise<TxEffect | undefined> {
+  getTxEffect(txHash: TxHash) {
     return this.pxe.getTxEffect(txHash);
   }
   getTxReceipt(txHash: TxHash): Promise<TxReceipt> {

yarn-project/circuit-types/src/in_block.ts

@@ -0,0 +1,36 @@
+import { Fr } from '@aztec/circuits.js';
+import { schemas } from '@aztec/foundation/schemas';
+
+import { type ZodTypeAny, z } from 'zod';
+
+import { type L2Block } from './l2_block.js';
+
+export type InBlock<T> = {
+  l2BlockNumber: number;
+  l2BlockHash: string;
+  data: T;
+};
+
+export function randomInBlock<T>(data: T): InBlock<T> {
+  return {
+    data,
+    l2BlockNumber: Math.floor(Math.random() * 1000),
+    l2BlockHash: Fr.random().toString(),
+  };
+}
+
+export function wrapInBlock<T>(data: T, block: L2Block): InBlock<T> {
+  return {
+    data,
+    l2BlockNumber: block.number,
+    l2BlockHash: block.hash().toString(),
+  };
+}
+
+export function inBlockSchemaFor<T extends ZodTypeAny>(schema: T) {
+  return z.object({
+    data: schema,
+    l2BlockNumber: schemas.Integer,
+    l2BlockHash: z.string(),
+  });
+}

yarn-project/circuit-types/src/index.ts

@@ -22,3 +22,4 @@ export * from './simulation_error.js';
 export * from './tx/index.js';
 export * from './tx_effect.js';
 export * from './tx_execution_request.js';
+export * from './in_block.js';

yarn-project/circuit-types/src/interfaces/archiver.test.ts

@@ -20,6 +20,7 @@ import { readFileSync } from 'fs';
 import omit from 'lodash.omit';
 import { resolve } from 'path';
 
+import { type InBlock } from '../in_block.js';
 import { L2Block } from '../l2_block.js';
 import { type L2Tips } from '../l2_block_source.js';
 import { ExtendedUnencryptedL2Log } from '../logs/extended_unencrypted_l2_log.js';
@@ -106,7 +107,7 @@ describe('ArchiverApiSchema', () => {
 
   it('getTxEffect', async () => {
     const result = await context.client.getTxEffect(new TxHash(Buffer.alloc(32, 1)));
-    expect(result).toBeInstanceOf(TxEffect);
+    expect(result!.data).toBeInstanceOf(TxEffect);
   });
 
   it('getSettledTxReceipt', async () => {
@@ -261,9 +262,9 @@ class MockArchiver implements ArchiverApi {
   getBlocks(from: number, _limit: number, _proven?: boolean | undefined): Promise<L2Block[]> {
     return Promise.resolve([L2Block.random(from)]);
   }
-  getTxEffect(_txHash: TxHash): Promise<TxEffect | undefined> {
+  getTxEffect(_txHash: TxHash): Promise<InBlock<TxEffect> | undefined> {
     expect(_txHash).toBeInstanceOf(TxHash);
-    return Promise.resolve(TxEffect.empty());
+    return Promise.resolve({ l2BlockNumber: 1, l2BlockHash: '0x12', data: TxEffect.random() });
   }
   getSettledTxReceipt(txHash: TxHash): Promise<TxReceipt | undefined> {
     expect(txHash).toBeInstanceOf(TxHash);