refactor(prover-client): deferred epoch finalization, checkpoint-ready promise, and reorg safety

yarn-project/end-to-end/bootstrap.sh

@@ -33,10 +33,12 @@ function test_cmds {
     echo "$prefix:NAME=e2e_prover_full_fake FAKE_PROOFS=1 $run_test_script simple e2e_prover/full"
   fi
   echo "$prefix:TIMEOUT=15m:NAME=e2e_block_building $(set_dump_avm e2e_block_building) $run_test_script simple e2e_block_building"
+  echo "$prefix:TIMEOUT=15m:NAME=e2e_epochs/epochs_long_proving_time $run_test_script simple src/e2e_epochs/epochs_long_proving_time.test.ts"
 
   local tests=(
     # List all standalone and nested tests, except for the ones listed above.
-    src/e2e_!(prover)/*.test.ts
+    src/e2e_!(prover|epochs)/*.test.ts
+    src/e2e_epochs/!(epochs_long_proving_time).test.ts
     src/e2e_p2p/reqresp/*.test.ts
     src/e2e_!(block_building).test.ts
   )

yarn-project/prover-client/src/mocks/test_context.ts

@@ -159,6 +159,29 @@ export class TestContext {
     this.epochNumber++;
   }
 
+  /** Removes the last checkpoint from the test context, rolling back state changes from makeCheckpoint. */
+  public async removeLastCheckpoint() {
+    const removed = this.checkpoints.pop();
+    this.checkpointOutHashes.pop();
+    if (!removed) {
+      return;
+    }
+
+    const numBlocks = removed.blocks.length;
+    this.nextCheckpointIndex--;
+    this.nextCheckpointNumber = CheckpointNumber(Number(this.nextCheckpointNumber) - 1);
+    this.nextBlockNumber -= numBlocks;
+
+    // Remove block headers.
+    for (const block of removed.blocks) {
+      this.headers.delete(block.number);
+    }
+
+    // Unwind world state to before the removed checkpoint's blocks.
+    const firstBlockNumber = removed.blocks[0].number;
+    await this.worldState.unwindBlocks(BlockNumber(firstBlockNumber - 1));
+  }
+
   // Return blob fields of all checkpoints in the epoch.
   public getBlobFields() {
     return this.checkpoints.map(checkpoint => checkpoint.toBlobFields());

yarn-project/prover-client/src/orchestrator/checkpoint-proving-state.ts

@@ -68,7 +68,7 @@ export class CheckpointProvingState {
     public readonly index: number,
     public readonly constants: CheckpointConstantData,
     public readonly totalNumBlocks: number,
-    private readonly finalBlobBatchingChallenges: FinalBlobBatchingChallenges,
+    private finalBlobBatchingChallenges: FinalBlobBatchingChallenges | undefined,
     private readonly headerOfLastBlockInPreviousCheckpoint: BlockHeader,
     private readonly lastArchiveSiblingPath: Tuple<Fr, typeof ARCHIVE_HEIGHT>,
     private readonly l1ToL2Messages: Fr[],
@@ -91,6 +91,22 @@ export class CheckpointProvingState {
     this.firstBlockNumber = BlockNumber(headerOfLastBlockInPreviousCheckpoint.globalVariables.blockNumber + 1);
   }
 
+  /** Sets the final blob batching challenges. Called from EpochProvingState.finalizeEpochStructure(). */
+  public setFinalBlobBatchingChallenges(challenges: FinalBlobBatchingChallenges) {
+    this.finalBlobBatchingChallenges = challenges;
+  }
+
+  /** Returns true if the block merge tree is fully resolved (all block root/merge proofs are ready). */
+  /** Returns true when all block-level proofs that feed into the checkpoint root are complete. */
+  public isBlockMergeTreeComplete(): boolean {
+    if (this.isAcceptingBlocks()) {
+      return false;
+    }
+    // Use the same check as isReadyForCheckpointRoot for the proof tree,
+    // but without requiring blob/out-hash data (which comes from epoch finalization).
+    return this.#getChildProofsForRoot().every(p => !!p);
+  }
+
   public get epochNumber(): number {
     return this.parentEpoch.epochNumber;
   }
@@ -283,6 +299,9 @@ export class CheckpointProvingState {
     if (!this.startBlobAccumulator) {
       throw new Error('Start blob accumulator is not set.');
     }
+    if (!this.finalBlobBatchingChallenges) {
+      throw new Error('Final blob batching challenges are not set. Call finalizeEpochStructure first.');
+    }
 
     // `blobFields` must've been set if `startBlobAccumulator` is set (in `accumulateBlobs`).
     const blobFields = this.blobFields!;