@aztec/prover-client 0.65.2 → 0.66.0

package/src/mocks/test_context.ts

@@ -1,6 +1,6 @@
 import { type BBProverConfig } from '@aztec/bb-prover';
 import {
-  type MerkleTreeWriteOperations,
+  type L2Block,
   type ProcessedTx,
   type ProcessedTxHandler,
   type PublicExecutionRequest,
@@ -8,10 +8,13 @@ import {
   type Tx,
   type TxValidator,
 } from '@aztec/circuit-types';
-import { type Gas, type GlobalVariables, Header } from '@aztec/circuits.js';
+import { makeBloatedProcessedTx } from '@aztec/circuit-types/test';
+import { type AppendOnlyTreeSnapshot, type Gas, type GlobalVariables, Header } from '@aztec/circuits.js';
+import { times } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/fields';
 import { type DebugLogger } from '@aztec/foundation/log';
-import { openTmpStore } from '@aztec/kv-store/utils';
+import { getVKTreeRoot } from '@aztec/noir-protocol-circuits-types';
+import { protocolContractTreeRoot } from '@aztec/protocol-contracts';
 import {
   PublicProcessor,
   PublicTxSimulator,
@@ -20,32 +23,34 @@ import {
   type WorldStateDB,
 } from '@aztec/simulator';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
-import { MerkleTrees } from '@aztec/world-state';
+import { type MerkleTreeAdminDatabase } from '@aztec/world-state';
 import { NativeWorldStateService } from '@aztec/world-state/native';
 
 import { jest } from '@jest/globals';
 import * as fs from 'fs/promises';
-import { type MockProxy, mock } from 'jest-mock-extended';
+import { mock } from 'jest-mock-extended';
 
 import { TestCircuitProver } from '../../../bb-prover/src/test/test_circuit_prover.js';
 import { AvmFinalizedCallResult } from '../../../simulator/src/avm/avm_contract_call_result.js';
 import { type AvmPersistableStateManager } from '../../../simulator/src/avm/journal/journal.js';
+import { buildBlock } from '../block_builder/light.js';
 import { ProvingOrchestrator } from '../orchestrator/index.js';
 import { MemoryProvingQueue } from '../prover-agent/memory-proving-queue.js';
 import { ProverAgent } from '../prover-agent/prover-agent.js';
 import { getEnvironmentConfig, getSimulationProvider, makeGlobals } from './fixtures.js';
 
 export class TestContext {
+  private headers: Map<number, Header> = new Map();
+
   constructor(
     public publicTxSimulator: PublicTxSimulator,
-    public worldStateDB: MockProxy<WorldStateDB>,
+    public worldState: MerkleTreeAdminDatabase,
     public publicProcessor: PublicProcessor,
     public simulationProvider: SimulationProvider,
     public globalVariables: GlobalVariables,
-    public actualDb: MerkleTreeWriteOperations,
     public prover: ServerCircuitProver,
     public proverAgent: ProverAgent,
-    public orchestrator: ProvingOrchestrator,
+    public orchestrator: TestProvingOrchestrator,
     public blockNumber: number,
     public directoriesToCleanup: string[],
     public logger: DebugLogger,
@@ -57,11 +62,10 @@ export class TestContext {
 
   static async new(
     logger: DebugLogger,
-    worldState: 'native' | 'legacy' = 'native',
     proverCount = 4,
     createProver: (bbConfig: BBProverConfig) => Promise<ServerCircuitProver> = _ =>
       Promise.resolve(new TestCircuitProver(new NoopTelemetryClient(), new WASMSimulator())),
-    blockNumber = 3,
+    blockNumber = 1,
   ) {
     const directoriesToCleanup: string[] = [];
     const globalVariables = makeGlobals(blockNumber);
@@ -70,18 +74,9 @@ export class TestContext {
     const telemetry = new NoopTelemetryClient();
 
     // Separated dbs for public processor and prover - see public_processor for context
-    let publicDb: MerkleTreeWriteOperations;
-    let proverDb: MerkleTreeWriteOperations;
+    const ws = await NativeWorldStateService.tmp();
+    const publicDb = await ws.fork();
 
-    if (worldState === 'native') {
-      const ws = await NativeWorldStateService.tmp();
-      publicDb = await ws.fork();
-      proverDb = await ws.fork();
-    } else {
-      const ws = await MerkleTrees.new(openTmpStore(), telemetry);
-      publicDb = await ws.getLatest();
-      proverDb = await ws.getLatest();
-    }
     worldStateDB.getMerkleInterface.mockReturnValue(publicDb);
 
     const publicTxSimulator = new PublicTxSimulator(publicDb, worldStateDB, telemetry, globalVariables);
@@ -118,7 +113,7 @@ export class TestContext {
     }
 
     const queue = new MemoryProvingQueue(telemetry);
-    const orchestrator = new ProvingOrchestrator(proverDb, queue, telemetry, Fr.ZERO);
+    const orchestrator = new TestProvingOrchestrator(ws, queue, telemetry, Fr.ZERO);
     const agent = new ProverAgent(localProver, proverCount);
 
     queue.start();
@@ -126,11 +121,10 @@ export class TestContext {
 
     return new this(
       publicTxSimulator,
-      worldStateDB,
+      ws,
       processor,
       simulationProvider,
       globalVariables,
-      proverDb,
       localProver,
       agent,
       orchestrator,
@@ -140,6 +134,16 @@ export class TestContext {
     );
   }
 
+  public getFork() {
+    return this.worldState.fork();
+  }
+
+  public getHeader(blockNumber: 0): Header;
+  public getHeader(blockNumber: number): Header | undefined;
+  public getHeader(blockNumber = 0) {
+    return blockNumber === 0 ? this.worldState.getCommitted().getInitialHeader() : this.headers.get(blockNumber);
+  }
+
   async cleanup() {
     await this.proverAgent.stop();
     for (const dir of this.directoriesToCleanup.filter(x => x !== '')) {
@@ -147,6 +151,42 @@ export class TestContext {
     }
   }
 
+  public makeProcessedTx(opts?: Parameters<typeof makeBloatedProcessedTx>[0]): ProcessedTx;
+  public makeProcessedTx(seed?: number): ProcessedTx;
+  public makeProcessedTx(seedOrOpts?: Parameters<typeof makeBloatedProcessedTx>[0] | number): ProcessedTx {
+    const opts = typeof seedOrOpts === 'number' ? { seed: seedOrOpts } : seedOrOpts;
+    const blockNum = (opts?.globalVariables ?? this.globalVariables).blockNumber.toNumber();
+    const header = this.getHeader(blockNum - 1);
+    return makeBloatedProcessedTx({
+      header,
+      vkTreeRoot: getVKTreeRoot(),
+      protocolContractTreeRoot,
+      globalVariables: this.globalVariables,
+      ...opts,
+    });
+  }
+
+  /** Creates a block with the given number of txs and adds it to world-state */
+  public async makePendingBlock(
+    numTxs: number,
+    numMsgs: number = 0,
+    blockNumOrGlobals: GlobalVariables | number = this.globalVariables,
+    makeProcessedTxOpts: (index: number) => Partial<Parameters<typeof makeBloatedProcessedTx>[0]> = () => ({}),
+  ) {
+    const globalVariables = typeof blockNumOrGlobals === 'number' ? makeGlobals(blockNumOrGlobals) : blockNumOrGlobals;
+    const blockNum = globalVariables.blockNumber.toNumber();
+    const db = await this.worldState.fork();
+    const msgs = times(numMsgs, i => new Fr(blockNum * 100 + i));
+    const txs = times(numTxs, i =>
+      this.makeProcessedTx({ seed: i + blockNum * 1000, globalVariables, ...makeProcessedTxOpts(i) }),
+    );
+
+    const block = await buildBlock(txs, globalVariables, msgs, db);
+    this.headers.set(blockNum, block.header);
+    await this.worldState.handleL2BlockAndMessages(block, msgs);
+    return { block, txs, msgs };
+  }
+
   public async processPublicFunctions(
     txs: Tx[],
     maxTransactions: number,
@@ -217,3 +257,19 @@ export class TestContext {
     return await this.publicProcessor.process(txs, maxTransactions, txHandler, txValidator);
   }
 }
+
+class TestProvingOrchestrator extends ProvingOrchestrator {
+  public isVerifyBuiltBlockAgainstSyncedStateEnabled = false;
+
+  // Disable this check by default, since it requires seeding world state with the block being built
+  // This is only enabled in some tests with multiple blocks that populate the pending chain via makePendingBlock
+  protected override verifyBuiltBlockAgainstSyncedState(
+    l2Block: L2Block,
+    newArchive: AppendOnlyTreeSnapshot,
+  ): Promise<void> {
+    if (this.isVerifyBuiltBlockAgainstSyncedStateEnabled) {
+      return super.verifyBuiltBlockAgainstSyncedState(l2Block, newArchive);
+    }
+    return Promise.resolve();
+  }
+}

package/src/orchestrator/epoch-proving-state.ts

@@ -50,20 +50,16 @@ export class EpochProvingState {
   private mergeRollupInputs: BlockMergeRollupInputData[] = [];
   public rootRollupPublicInputs: RootRollupPublicInputs | undefined;
   public finalProof: Proof | undefined;
-  public blocks: BlockProvingState[] = [];
+  public blocks: (BlockProvingState | undefined)[] = [];
 
   constructor(
     public readonly epochNumber: number,
+    public readonly firstBlockNumber: number,
     public readonly totalNumBlocks: number,
     private completionCallback: (result: ProvingResult) => void,
     private rejectionCallback: (reason: string) => void,
   ) {}
 
-  /** Returns the current block proving state */
-  public get currentBlock(): BlockProvingState | undefined {
-    return this.blocks.at(-1);
-  }
-
   // Returns the number of levels of merge rollups
   public get numMergeLevels() {
     const totalLeaves = Math.max(2, this.totalNumBlocks);
@@ -110,9 +106,10 @@ export class EpochProvingState {
     archiveTreeSnapshot: AppendOnlyTreeSnapshot,
     archiveTreeRootSiblingPath: Tuple<Fr, typeof ARCHIVE_HEIGHT>,
     previousBlockHash: Fr,
-  ) {
+  ): BlockProvingState {
+    const index = globalVariables.blockNumber.toNumber() - this.firstBlockNumber;
     const block = new BlockProvingState(
-      this.blocks.length,
+      index,
       numTxs,
       globalVariables,
       padArrayEnd(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP),
@@ -124,11 +121,11 @@ export class EpochProvingState {
       previousBlockHash,
       this,
     );
-    this.blocks.push(block);
-    if (this.blocks.length === this.totalNumBlocks) {
+    this.blocks[index] = block;
+    if (this.blocks.filter(b => !!b).length === this.totalNumBlocks) {
       this.provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_FULL;
     }
-    return this.blocks.length - 1;
+    return block;
   }
 
   // Returns true if this proving state is still valid, false otherwise
@@ -180,8 +177,8 @@ export class EpochProvingState {
   }
 
   // Returns a specific transaction proving state
-  public getBlockProvingState(index: number) {
-    return this.blocks[index];
+  public getBlockProvingStateByBlockNumber(blockNumber: number) {
+    return this.blocks.find(block => block?.blockNumber === blockNumber);
   }
 
   // Returns a set of merge rollup inputs

package/src/orchestrator/orchestrator.ts

@@ -7,6 +7,7 @@ import {
 } from '@aztec/circuit-types';
 import {
   type EpochProver,
+  type ForkMerkleTreeOperations,
   type MerkleTreeWriteOperations,
   type ProofAndVerificationKey,
 } from '@aztec/circuit-types/interfaces';
@@ -14,6 +15,7 @@ import { type CircuitName } from '@aztec/circuit-types/stats';
 import {
   AVM_PROOF_LENGTH_IN_FIELDS,
   AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS,
+  type AppendOnlyTreeSnapshot,
   type BaseOrMergeRollupPublicInputs,
   BaseParityInputs,
   type BaseRollupHints,
@@ -38,7 +40,7 @@ import {
   makeEmptyRecursiveProof,
 } from '@aztec/circuits.js';
 import { makeTuple } from '@aztec/foundation/array';
-import { padArrayEnd } from '@aztec/foundation/collection';
+import { maxBy, padArrayEnd } from '@aztec/foundation/collection';
 import { AbortError } from '@aztec/foundation/error';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { promiseWithResolvers } from '@aztec/foundation/promise';
@@ -98,9 +100,10 @@ export class ProvingOrchestrator implements EpochProver {
 
   private provingPromise: Promise<ProvingResult> | undefined = undefined;
   private metrics: ProvingOrchestratorMetrics;
+  private dbs: Map<number, MerkleTreeWriteOperations> = new Map();
 
   constructor(
-    private db: MerkleTreeWriteOperations,
+    private dbProvider: ForkMerkleTreeOperations,
     private prover: ServerCircuitProver,
     telemetryClient: TelemetryClient,
     private readonly proverId: Fr = Fr.ZERO,
@@ -123,14 +126,14 @@ export class ProvingOrchestrator implements EpochProver {
     this.paddingTxProof = undefined;
   }
 
-  public startNewEpoch(epochNumber: number, totalNumBlocks: number) {
+  public startNewEpoch(epochNumber: number, firstBlockNumber: number, totalNumBlocks: number) {
     const { promise: _promise, resolve, reject } = promiseWithResolvers<ProvingResult>();
     const promise = _promise.catch((reason): ProvingResult => ({ status: 'failure', reason }));
     if (totalNumBlocks <= 0 || !Number.isInteger(totalNumBlocks)) {
       throw new Error(`Invalid number of blocks for epoch (got ${totalNumBlocks})`);
     }
     logger.info(`Starting epoch ${epochNumber} with ${totalNumBlocks} blocks`);
-    this.provingState = new EpochProvingState(epochNumber, totalNumBlocks, resolve, reject);
+    this.provingState = new EpochProvingState(epochNumber, firstBlockNumber, totalNumBlocks, resolve, reject);
     this.provingPromise = promise;
   }
 
@@ -159,24 +162,14 @@ export class ProvingOrchestrator implements EpochProver {
       throw new Error(`Invalid number of txs for block (got ${numTxs})`);
     }
 
-    if (this.provingState.currentBlock && !this.provingState.currentBlock.block) {
-      throw new Error(`Must end previous block before starting a new one`);
-    }
-
-    // TODO(palla/prover): Store block number in the db itself to make this check more reliable,
-    // and turn this warning into an exception that we throw.
-    const { blockNumber } = globalVariables;
-    const dbBlockNumber = (await this.db.getTreeInfo(MerkleTreeId.ARCHIVE)).size - 1n;
-    if (dbBlockNumber !== blockNumber.toBigInt() - 1n) {
-      logger.warn(
-        `Database is at wrong block number (starting block ${blockNumber.toBigInt()} with db at ${dbBlockNumber})`,
-      );
-    }
-
     logger.info(
-      `Starting block ${globalVariables.blockNumber} for slot ${globalVariables.slotNumber} with ${numTxs} transactions`,
+      `Starting block ${globalVariables.blockNumber.toNumber()} for slot ${globalVariables.slotNumber.toNumber()} with ${numTxs} transactions`,
     );
 
+    // Fork world state at the end of the immediately previous block
+    const db = await this.dbProvider.fork(globalVariables.blockNumber.toNumber() - 1);
+    this.dbs.set(globalVariables.blockNumber.toNumber(), db);
+
     // we start the block by enqueueing all of the base parity circuits
     let baseParityInputs: BaseParityInputs[] = [];
     let l1ToL2MessagesPadded: Tuple<Fr, typeof NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP>;
@@ -189,12 +182,12 @@ export class ProvingOrchestrator implements EpochProver {
       BaseParityInputs.fromSlice(l1ToL2MessagesPadded, i, getVKTreeRoot()),
     );
 
-    const messageTreeSnapshot = await getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, this.db);
+    const messageTreeSnapshot = await getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, db);
 
     const newL1ToL2MessageTreeRootSiblingPathArray = await getSubtreeSiblingPath(
       MerkleTreeId.L1_TO_L2_MESSAGE_TREE,
       L1_TO_L2_MSG_SUBTREE_HEIGHT,
-      this.db,
+      db,
     );
 
     const newL1ToL2MessageTreeRootSiblingPath = makeTuple(
@@ -205,18 +198,18 @@ export class ProvingOrchestrator implements EpochProver {
     );
 
     // Update the local trees to include the new l1 to l2 messages
-    await this.db.appendLeaves(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l1ToL2MessagesPadded);
-    const messageTreeSnapshotAfterInsertion = await getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, this.db);
+    await db.appendLeaves(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l1ToL2MessagesPadded);
+    const messageTreeSnapshotAfterInsertion = await getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, db);
 
     // Get archive snapshot before this block lands
-    const startArchiveSnapshot = await getTreeSnapshot(MerkleTreeId.ARCHIVE, this.db);
-    const newArchiveSiblingPath = await getRootTreeSiblingPath(MerkleTreeId.ARCHIVE, this.db);
-    const previousBlockHash = await this.db.getLeafValue(
+    const startArchiveSnapshot = await getTreeSnapshot(MerkleTreeId.ARCHIVE, db);
+    const newArchiveSiblingPath = await getRootTreeSiblingPath(MerkleTreeId.ARCHIVE, db);
+    const previousBlockHash = await db.getLeafValue(
       MerkleTreeId.ARCHIVE,
       BigInt(startArchiveSnapshot.nextAvailableLeafIndex - 1),
     );
 
-    this.provingState!.startNewBlock(
+    const blockProvingState = this.provingState!.startNewBlock(
       numTxs,
       globalVariables,
       l1ToL2MessagesPadded,
@@ -230,7 +223,7 @@ export class ProvingOrchestrator implements EpochProver {
 
     // Enqueue base parity circuits for the block
     for (let i = 0; i < baseParityInputs.length; i++) {
-      this.enqueueBaseParityCircuit(this.provingState!.currentBlock!, baseParityInputs[i], i);
+      this.enqueueBaseParityCircuit(blockProvingState, baseParityInputs[i], i);
     }
   }
 
@@ -242,33 +235,40 @@ export class ProvingOrchestrator implements EpochProver {
     [Attributes.TX_HASH]: tx.hash.toString(),
   }))
   public async addNewTx(tx: ProcessedTx): Promise<void> {
-    const provingState = this?.provingState?.currentBlock;
-    if (!provingState) {
-      throw new Error(`Invalid proving state, call startNewBlock before adding transactions`);
-    }
+    const blockNumber = tx.constants.globalVariables.blockNumber.toNumber();
+    try {
+      const provingState = this.provingState?.getBlockProvingStateByBlockNumber(blockNumber);
+      if (!provingState) {
+        throw new Error(`Block proving state for ${blockNumber} not found`);
+      }
 
-    if (!provingState.isAcceptingTransactions()) {
-      throw new Error(`Rollup not accepting further transactions`);
-    }
+      if (!provingState.isAcceptingTransactions()) {
+        throw new Error(`Rollup not accepting further transactions`);
+      }
 
-    if (!provingState.verifyState()) {
-      throw new Error(`Invalid proving state when adding a tx`);
-    }
+      if (!provingState.verifyState()) {
+        throw new Error(`Invalid proving state when adding a tx`);
+      }
 
-    validateTx(tx);
+      validateTx(tx);
 
-    logger.info(`Received transaction: ${tx.hash}`);
+      logger.info(`Received transaction: ${tx.hash}`);
 
-    if (tx.isEmpty) {
-      logger.warn(`Ignoring empty transaction ${tx.hash} - it will not be added to this block`);
-      return;
-    }
+      if (tx.isEmpty) {
+        logger.warn(`Ignoring empty transaction ${tx.hash} - it will not be added to this block`);
+        return;
+      }
 
-    const [hints, treeSnapshots] = await this.prepareTransaction(tx, provingState);
-    this.enqueueFirstProofs(hints, treeSnapshots, tx, provingState);
+      const [hints, treeSnapshots] = await this.prepareTransaction(tx, provingState);
+      this.enqueueFirstProofs(hints, treeSnapshots, tx, provingState);
 
-    if (provingState.transactionsReceived === provingState.totalNumTxs) {
-      logger.verbose(`All transactions received for block ${provingState.globalVariables.blockNumber}.`);
+      if (provingState.transactionsReceived === provingState.totalNumTxs) {
+        logger.verbose(`All transactions received for block ${provingState.globalVariables.blockNumber}.`);
+      }
+    } catch (err: any) {
+      throw new Error(`Error adding transaction ${tx.hash.toString()} to block ${blockNumber}: ${err.message}`, {
+        cause: err,
+      });
     }
   }
 
@@ -276,21 +276,13 @@ export class ProvingOrchestrator implements EpochProver {
    * Marks the block as full and pads it if required, no more transactions will be accepted.
    * Computes the block header and updates the archive tree.
    */
-  @trackSpan('ProvingOrchestrator.setBlockCompleted', function () {
-    const block = this.provingState?.currentBlock;
-    if (!block) {
-      return {};
-    }
-    return {
-      [Attributes.BLOCK_NUMBER]: block.globalVariables.blockNumber.toNumber(),
-      [Attributes.BLOCK_SIZE]: block.totalNumTxs,
-      [Attributes.BLOCK_TXS_COUNT]: block.transactionsReceived,
-    };
-  })
-  public async setBlockCompleted(expectedHeader?: Header): Promise<L2Block> {
-    const provingState = this.provingState?.currentBlock;
+  @trackSpan('ProvingOrchestrator.setBlockCompleted', (blockNumber: number) => ({
+    [Attributes.BLOCK_NUMBER]: blockNumber,
+  }))
+  public async setBlockCompleted(blockNumber: number, expectedHeader?: Header): Promise<L2Block> {
+    const provingState = this.provingState?.getBlockProvingStateByBlockNumber(blockNumber);
     if (!provingState) {
-      throw new Error(`Invalid proving state, call startNewBlock before adding transactions or completing the block`);
+      throw new Error(`Block proving state for ${blockNumber} not found`);
     }
 
     if (!provingState.verifyState()) {
@@ -313,7 +305,7 @@ export class ProvingOrchestrator implements EpochProver {
       // base rollup inputs
       // Then enqueue the proving of all the transactions
       const unprovenPaddingTx = makeEmptyProcessedTx(
-        this.db.getInitialHeader(),
+        this.dbs.get(blockNumber)!.getInitialHeader(),
         provingState.globalVariables.chainId,
         provingState.globalVariables.version,
         getVKTreeRoot(),
@@ -344,7 +336,7 @@ export class ProvingOrchestrator implements EpochProver {
 
   /** Returns the block as built for a given index. */
   public getBlock(index: number): L2Block {
-    const block = this.provingState?.blocks[index].block;
+    const block = this.provingState?.blocks[index]?.block;
     if (!block) {
       throw new Error(`Block at index ${index} not available`);
     }
@@ -362,7 +354,10 @@ export class ProvingOrchestrator implements EpochProver {
   })
   private padEpoch(): Promise<void> {
     const provingState = this.provingState!;
-    const lastBlock = provingState.currentBlock?.block;
+    const lastBlock = maxBy(
+      provingState.blocks.filter(b => !!b),
+      b => b!.blockNumber,
+    )?.block;
     if (!lastBlock) {
       return Promise.reject(new Error(`Epoch needs at least one completed block in order to be padded`));
     }
@@ -416,13 +411,16 @@ export class ProvingOrchestrator implements EpochProver {
     // Collect all new nullifiers, commitments, and contracts from all txs in this block to build body
     const txs = provingState!.allTxs.map(a => a.processedTx);
 
+    // Get db for this block
+    const db = this.dbs.get(provingState.blockNumber)!;
+
     // Given we've applied every change from this block, now assemble the block header
     // and update the archive tree, so we're ready to start processing the next block
     const { header, body } = await buildHeaderAndBodyFromTxs(
       txs,
       provingState.globalVariables,
       provingState.newL1ToL2Messages,
-      this.db,
+      db,
     );
 
     if (expectedHeader && !header.equals(expectedHeader)) {
@@ -431,10 +429,10 @@ export class ProvingOrchestrator implements EpochProver {
     }
 
     logger.verbose(`Updating archive tree with block ${provingState.blockNumber} header ${header.hash().toString()}`);
-    await this.db.updateArchive(header);
+    await db.updateArchive(header);
 
     // Assemble the L2 block
-    const newArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, this.db);
+    const newArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, db);
     const l2Block = new L2Block(newArchive, header, body);
 
     if (!l2Block.body.getTxsEffectsHash().equals(header.contentCommitment.txsEffectsHash)) {
@@ -445,10 +443,24 @@ export class ProvingOrchestrator implements EpochProver {
       );
     }
 
+    await this.verifyBuiltBlockAgainstSyncedState(l2Block, newArchive);
+
     logger.verbose(`Orchestrator finalised block ${l2Block.number}`);
     provingState.block = l2Block;
   }
 
+  // Flagged as protected to disable in certain unit tests
+  protected async verifyBuiltBlockAgainstSyncedState(l2Block: L2Block, newArchive: AppendOnlyTreeSnapshot) {
+    const syncedArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, this.dbProvider.getSnapshot(l2Block.number));
+    if (!syncedArchive.equals(newArchive)) {
+      throw new Error(
+        `Archive tree mismatch for block ${l2Block.number}: world state synced to ${inspect(
+          syncedArchive,
+        )} but built ${inspect(newArchive)}`,
+      );
+    }
+  }
+
   // Enqueues the proving of the required padding transactions
   // If the fully proven padding transaction is not available, this will first be proven
   private enqueuePaddingTxs(
@@ -602,13 +614,6 @@ export class ProvingOrchestrator implements EpochProver {
     provingState: BlockProvingState,
   ) {
     const txProvingState = new TxProvingState(tx, hints, treeSnapshots);
-
-    const rejectReason = txProvingState.verifyStateOrReject();
-    if (rejectReason) {
-      provingState.reject(rejectReason);
-      return;
-    }
-
     const txIndex = provingState.addNewTx(txProvingState);
     this.enqueueTube(provingState, txIndex);
     if (txProvingState.requireAvmProof) {
@@ -692,9 +697,11 @@ export class ProvingOrchestrator implements EpochProver {
       return;
     }
 
+    const db = this.dbs.get(provingState.blockNumber)!;
+
     // We build the base rollup inputs using a mock proof and verification key.
     // These will be overwritten later once we have proven the tube circuit and any public kernels
-    const [ms, hints] = await elapsed(buildBaseRollupHints(tx, provingState.globalVariables, this.db));
+    const [ms, hints] = await elapsed(buildBaseRollupHints(tx, provingState.globalVariables, db));
 
     if (!tx.isEmpty) {
       this.metrics.recordBaseRollupInputs(ms);
@@ -702,7 +709,7 @@ export class ProvingOrchestrator implements EpochProver {
 
     const promises = [MerkleTreeId.NOTE_HASH_TREE, MerkleTreeId.NULLIFIER_TREE, MerkleTreeId.PUBLIC_DATA_TREE].map(
       async (id: MerkleTreeId) => {
-        return { key: id, value: await getTreeSnapshot(id, this.db) };
+        return { key: id, value: await getTreeSnapshot(id, db) };
       },
     );
     const treeSnapshots: TreeSnapshots = new Map((await Promise.all(promises)).map(obj => [obj.key, obj.value]));
@@ -1055,6 +1062,19 @@ export class ProvingOrchestrator implements EpochProver {
       logger.debug('Block root rollup already started');
       return;
     }
+    const blockNumber = provingState.blockNumber;
+
+    // TODO(palla/prover): This closes the fork only on the happy path. If this epoch orchestrator
+    // is aborted and never reaches this point, it will leak the fork. We need to add a global cleanup,
+    // but have to make sure it only runs once all operations are completed, otherwise some function here
+    // will attempt to access the fork after it was closed.
+    logger.debug(`Cleaning up world state fork for ${blockNumber}`);
+    void this.dbs
+      .get(blockNumber)
+      ?.close()
+      .then(() => this.dbs.delete(blockNumber))
+      .catch(err => logger.error(`Error closing db for block ${blockNumber}`, err));
+
     this.enqueueBlockRootRollup(provingState);
   }