@aztec/prover-client 0.0.0-test.0 → 0.0.1-commit.03f7ef2

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

@@ -1,27 +1,31 @@
-import {
-  type ARCHIVE_HEIGHT,
-  type L1_TO_L2_MSG_SUBTREE_SIBLING_PATH_LENGTH,
-  type NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH,
-  type TUBE_PROOF_LENGTH,
-  VK_TREE_HEIGHT,
+import { BatchedBlob, BatchedBlobAccumulator, type FinalBlobBatchingChallenges } from '@aztec/blob-lib';
+import type {
+  ARCHIVE_HEIGHT,
+  L1_TO_L2_MSG_SUBTREE_ROOT_SIBLING_PATH_LENGTH,
+  NESTED_RECURSIVE_PROOF_LENGTH,
+  NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH,
 } from '@aztec/constants';
-import type { Fr } from '@aztec/foundation/fields';
+import { BlockNumber, EpochNumber } from '@aztec/foundation/branded-types';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { Tuple } from '@aztec/foundation/serialize';
-import { MembershipWitness, type TreeNodeLocation, UnbalancedTreeStore } from '@aztec/foundation/trees';
-import { getVKIndex, getVKSiblingPath } from '@aztec/noir-protocol-circuits-types/vk-tree';
-import type { ProofAndVerificationKey, PublicInputsAndRecursiveProof } from '@aztec/stdlib/interfaces/server';
+import { type TreeNodeLocation, UnbalancedTreeStore } from '@aztec/foundation/trees';
+import type { PublicInputsAndRecursiveProof } from '@aztec/stdlib/interfaces/server';
 import type { Proof } from '@aztec/stdlib/proofs';
 import {
-  BlockMergeRollupInputs,
-  type BlockRootOrBlockMergePublicInputs,
-  PreviousRollupBlockData,
-  RootRollupInputs,
+  CheckpointConstantData,
+  CheckpointMergeRollupPrivateInputs,
+  CheckpointPaddingRollupPrivateInputs,
+  CheckpointRollupPublicInputs,
+  PublicChonkVerifierPublicInputs,
+  RootRollupPrivateInputs,
   type RootRollupPublicInputs,
 } from '@aztec/stdlib/rollup';
 import type { AppendOnlyTreeSnapshot, MerkleTreeId } from '@aztec/stdlib/trees';
-import type { BlockHeader, GlobalVariables } from '@aztec/stdlib/tx';
+import type { BlockHeader } from '@aztec/stdlib/tx';
 
-import { BlockProvingState } from './block-proving-state.js';
+import { toProofData } from './block-building-helpers.js';
+import type { ProofState } from './block-proving-state.js';
+import { CheckpointProvingState } from './checkpoint-proving-state.js';
 
 export type TreeSnapshots = Map<MerkleTreeId, AppendOnlyTreeSnapshot>;
 
@@ -41,58 +45,98 @@ export type ProvingResult = { status: 'success' } | { status: 'failure'; reason:
  * Captures resolve and reject callbacks to provide a promise base interface to the consumer of our proving.
  */
 export class EpochProvingState {
-  private blockRootOrMergeProvingOutputs: UnbalancedTreeStore<
-    PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+  private checkpointProofs: UnbalancedTreeStore<
+    ProofState<CheckpointRollupPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
   >;
-  private paddingBlockRootProvingOutput:
-    | PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+  private checkpointPaddingProof:
+    | ProofState<CheckpointRollupPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
     | undefined;
-  private rootRollupProvingOutput: PublicInputsAndRecursiveProof<RootRollupPublicInputs> | undefined;
+  private rootRollupProof: ProofState<RootRollupPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined;
+  private checkpoints: (CheckpointProvingState | undefined)[] = [];
+  private startBlobAccumulator: BatchedBlobAccumulator;
+  private endBlobAccumulator: BatchedBlobAccumulator | undefined;
+  private finalBatchedBlob: BatchedBlob | undefined;
   private provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED;
 
-  // Map from tx hash to tube proof promise. Used when kickstarting tube proofs before tx processing.
-  public readonly cachedTubeProofs = new Map<string, Promise<ProofAndVerificationKey<typeof TUBE_PROOF_LENGTH>>>();
-
-  public blocks: (BlockProvingState | undefined)[] = [];
+  // Map from tx hash to chonk verifier proof promise. Used when kickstarting chonk verifier proofs before tx processing.
+  public readonly cachedChonkVerifierProofs = new Map<
+    string,
+    Promise<
+      PublicInputsAndRecursiveProof<PublicChonkVerifierPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+    >
+  >();
 
   constructor(
-    public readonly epochNumber: number,
-    public readonly firstBlockNumber: number,
-    public readonly totalNumBlocks: number,
+    public readonly epochNumber: EpochNumber,
+    public readonly totalNumCheckpoints: number,
+    private readonly finalBlobBatchingChallenges: FinalBlobBatchingChallenges,
+    private onCheckpointBlobAccumulatorSet: (checkpoint: CheckpointProvingState) => void,
     private completionCallback: (result: ProvingResult) => void,
     private rejectionCallback: (reason: string) => void,
   ) {
-    this.blockRootOrMergeProvingOutputs = new UnbalancedTreeStore(totalNumBlocks);
+    this.checkpointProofs = new UnbalancedTreeStore(totalNumCheckpoints);
+    this.startBlobAccumulator = BatchedBlobAccumulator.newWithChallenges(finalBlobBatchingChallenges);
   }
 
   // Adds a block to the proving state, returns its index
   // Will update the proving life cycle if this is the last block
-  public startNewBlock(
-    globalVariables: GlobalVariables,
-    l1ToL2Messages: Fr[],
-    l1ToL2MessageSubtreeSiblingPath: Tuple<Fr, typeof L1_TO_L2_MSG_SUBTREE_SIBLING_PATH_LENGTH>,
-    l1ToL2MessageTreeSnapshotAfterInsertion: AppendOnlyTreeSnapshot,
-    lastArchiveSnapshot: AppendOnlyTreeSnapshot,
-    newArchiveSiblingPath: Tuple<Fr, typeof ARCHIVE_HEIGHT>,
+  public startNewCheckpoint(
+    checkpointIndex: number,
+    constants: CheckpointConstantData,
+    totalNumBlocks: number,
     previousBlockHeader: BlockHeader,
-  ): BlockProvingState {
-    const index = globalVariables.blockNumber.toNumber() - this.firstBlockNumber;
-    const block = new BlockProvingState(
-      index,
-      globalVariables,
-      l1ToL2Messages,
-      l1ToL2MessageSubtreeSiblingPath,
-      l1ToL2MessageTreeSnapshotAfterInsertion,
-      lastArchiveSnapshot,
-      newArchiveSiblingPath,
+    lastArchiveSiblingPath: Tuple<Fr, typeof ARCHIVE_HEIGHT>,
+    l1ToL2Messages: Fr[],
+    lastL1ToL2MessageTreeSnapshot: AppendOnlyTreeSnapshot,
+    lastL1ToL2MessageSubtreeRootSiblingPath: Tuple<Fr, typeof L1_TO_L2_MSG_SUBTREE_ROOT_SIBLING_PATH_LENGTH>,
+    newL1ToL2MessageTreeSnapshot: AppendOnlyTreeSnapshot,
+    newL1ToL2MessageSubtreeRootSiblingPath: Tuple<Fr, typeof L1_TO_L2_MSG_SUBTREE_ROOT_SIBLING_PATH_LENGTH>,
+  ): CheckpointProvingState {
+    if (checkpointIndex >= this.totalNumCheckpoints) {
+      throw new Error(
+        `Unable to start a new checkpoint at index ${checkpointIndex}. Expected at most ${this.totalNumCheckpoints} checkpoints.`,
+      );
+    }
+
+    const checkpoint = new CheckpointProvingState(
+      checkpointIndex,
+      constants,
+      totalNumBlocks,
+      this.finalBlobBatchingChallenges,
       previousBlockHeader,
+      lastArchiveSiblingPath,
+      l1ToL2Messages,
+      lastL1ToL2MessageTreeSnapshot,
+      lastL1ToL2MessageSubtreeRootSiblingPath,
+      newL1ToL2MessageTreeSnapshot,
+      newL1ToL2MessageSubtreeRootSiblingPath,
       this,
+      this.onCheckpointBlobAccumulatorSet,
     );
-    this.blocks[index] = block;
-    if (this.blocks.filter(b => !!b).length === this.totalNumBlocks) {
+    this.checkpoints[checkpointIndex] = checkpoint;
+
+    if (this.checkpoints.filter(c => !!c).length === this.totalNumCheckpoints) {
       this.provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_FULL;
     }
-    return block;
+
+    return checkpoint;
+  }
+
+  public getCheckpointProvingState(index: number) {
+    return this.checkpoints[index];
+  }
+
+  public getCheckpointProvingStateByBlockNumber(blockNumber: BlockNumber) {
+    return this.checkpoints.find(
+      c =>
+        c &&
+        Number(blockNumber) >= Number(c.firstBlockNumber) &&
+        Number(blockNumber) < Number(c.firstBlockNumber) + c.totalNumBlocks,
+    );
+  }
+
+  public getBlockProvingStateByBlockNumber(blockNumber: BlockNumber) {
+    return this.getCheckpointProvingStateByBlockNumber(blockNumber)?.getBlockProvingStateByBlockNumber(blockNumber);
   }
 
   // Returns true if this proving state is still valid, false otherwise
@@ -103,95 +147,146 @@ export class EpochProvingState {
     );
   }
 
-  // Returns true if we are still able to accept blocks, false otherwise
-  public isAcceptingBlocks() {
-    return this.provingStateLifecycle === PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED;
+  // Returns true if we are still able to accept checkpoints, false otherwise.
+  public isAcceptingCheckpoints() {
+    return this.checkpoints.filter(c => !!c).length < this.totalNumCheckpoints;
   }
 
-  public setBlockRootRollupProof(
-    blockIndex: number,
-    proof: PublicInputsAndRecursiveProof<
-      BlockRootOrBlockMergePublicInputs,
+  public setCheckpointRootRollupProof(
+    checkpointIndex: number,
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
       typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
     >,
   ): TreeNodeLocation {
-    return this.blockRootOrMergeProvingOutputs.setLeaf(blockIndex, proof);
+    return this.checkpointProofs.setLeaf(checkpointIndex, { provingOutput });
+  }
+
+  public tryStartProvingCheckpointMerge(location: TreeNodeLocation) {
+    if (this.checkpointProofs.getNode(location)?.isProving) {
+      return false;
+    } else {
+      this.checkpointProofs.setNode(location, { isProving: true });
+      return true;
+    }
   }
 
-  public setBlockMergeRollupProof(
+  public setCheckpointMergeRollupProof(
     location: TreeNodeLocation,
-    proof: PublicInputsAndRecursiveProof<
-      BlockRootOrBlockMergePublicInputs,
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
       typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
     >,
   ) {
-    this.blockRootOrMergeProvingOutputs.setNode(location, proof);
+    this.checkpointProofs.setNode(location, { provingOutput });
+  }
+
+  public tryStartProvingRootRollup() {
+    if (this.rootRollupProof?.isProving) {
+      return false;
+    } else {
+      this.rootRollupProof = { isProving: true };
+      return true;
+    }
   }
 
-  public setRootRollupProof(proof: PublicInputsAndRecursiveProof<RootRollupPublicInputs>) {
-    this.rootRollupProvingOutput = proof;
+  public setRootRollupProof(provingOutput: PublicInputsAndRecursiveProof<RootRollupPublicInputs>) {
+    this.rootRollupProof = { provingOutput };
   }
 
-  public setPaddingBlockRootProof(
-    proof: PublicInputsAndRecursiveProof<
-      BlockRootOrBlockMergePublicInputs,
+  public tryStartProvingPaddingCheckpoint() {
+    if (this.checkpointPaddingProof?.isProving) {
+      return false;
+    } else {
+      this.checkpointPaddingProof = { isProving: true };
+      return true;
+    }
+  }
+
+  public setCheckpointPaddingProof(
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
       typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
     >,
   ) {
-    this.paddingBlockRootProvingOutput = proof;
+    this.checkpointPaddingProof = { provingOutput };
+  }
+
+  public async setBlobAccumulators() {
+    let previousAccumulator = this.startBlobAccumulator;
+    // Accumulate blobs as far as we can for this epoch.
+    for (let i = 0; i < this.totalNumCheckpoints; i++) {
+      const checkpoint = this.checkpoints[i];
+      if (!checkpoint) {
+        break;
+      }
+
+      const endAccumulator =
+        checkpoint.getEndBlobAccumulator() || (await checkpoint.accumulateBlobs(previousAccumulator));
+      if (!endAccumulator) {
+        break;
+      }
+
+      previousAccumulator = endAccumulator;
+
+      // If this is the last checkpoint, set the end blob accumulator.
+      if (i === this.totalNumCheckpoints - 1) {
+        this.endBlobAccumulator = endAccumulator;
+      }
+    }
+  }
+
+  public async finalizeBatchedBlob() {
+    if (!this.endBlobAccumulator) {
+      throw new Error('End blob accumulator not ready.');
+    }
+    this.finalBatchedBlob = await this.endBlobAccumulator.finalize(true /* verifyProof */);
   }
 
   public getParentLocation(location: TreeNodeLocation) {
-    return this.blockRootOrMergeProvingOutputs.getParentLocation(location);
+    return this.checkpointProofs.getParentLocation(location);
   }
 
-  public getBlockMergeRollupInputs(mergeLocation: TreeNodeLocation) {
-    const [left, right] = this.blockRootOrMergeProvingOutputs.getChildren(mergeLocation);
+  public getCheckpointMergeRollupInputs(mergeLocation: TreeNodeLocation) {
+    const [left, right] = this.checkpointProofs.getChildren(mergeLocation).map(c => c?.provingOutput);
     if (!left || !right) {
-      throw new Error('At lease one child is not ready.');
+      throw new Error('At least one child is not ready for the checkpoint merge rollup.');
     }
 
-    return new BlockMergeRollupInputs([this.#getPreviousRollupData(left), this.#getPreviousRollupData(right)]);
+    return new CheckpointMergeRollupPrivateInputs([toProofData(left), toProofData(right)]);
   }
 
-  public getRootRollupInputs(proverId: Fr) {
+  public getRootRollupInputs() {
     const [left, right] = this.#getChildProofsForRoot();
     if (!left || !right) {
-      throw new Error('At lease one child is not ready.');
+      throw new Error('At least one child is not ready for the root rollup.');
     }
 
-    return RootRollupInputs.from({
-      previousRollupData: [this.#getPreviousRollupData(left), this.#getPreviousRollupData(right)],
-      proverId,
+    return RootRollupPrivateInputs.from({
+      previousRollups: [toProofData(left), toProofData(right)],
     });
   }
 
-  public getPaddingBlockRootInputs(proverId: Fr) {
-    if (!this.blocks[0]?.isComplete()) {
-      throw new Error('Epoch needs one completed block in order to be padded.');
-    }
-
-    return this.blocks[0].getPaddingBlockRootInputs(proverId);
+  public getPaddingCheckpointInputs() {
+    return new CheckpointPaddingRollupPrivateInputs();
   }
 
-  // Returns a specific transaction proving state
-  public getBlockProvingStateByBlockNumber(blockNumber: number) {
-    return this.blocks.find(block => block?.blockNumber === blockNumber);
-  }
+  public getEpochProofResult(): { proof: Proof; publicInputs: RootRollupPublicInputs; batchedBlobInputs: BatchedBlob } {
+    const provingOutput = this.rootRollupProof?.provingOutput;
 
-  public getEpochProofResult(): { proof: Proof; publicInputs: RootRollupPublicInputs } {
-    if (!this.rootRollupProvingOutput) {
+    if (!provingOutput || !this.finalBatchedBlob) {
       throw new Error('Unable to get epoch proof result. Root rollup is not ready.');
     }
 
     return {
-      proof: this.rootRollupProvingOutput.proof.binaryProof,
-      publicInputs: this.rootRollupProvingOutput.inputs,
+      proof: provingOutput.proof.binaryProof,
+      publicInputs: provingOutput.inputs,
+      batchedBlobInputs: this.finalBatchedBlob,
     };
   }
 
-  public isReadyForBlockMerge(location: TreeNodeLocation) {
-    return this.blockRootOrMergeProvingOutputs.getSibling(location) !== undefined;
+  public isReadyForCheckpointMerge(location: TreeNodeLocation) {
+    return !!this.checkpointProofs.getSibling(location)?.provingOutput;
   }
 
   // Returns true if we have sufficient inputs to execute the block root rollup
@@ -228,25 +323,8 @@ export class EpochProvingState {
   #getChildProofsForRoot() {
     const rootLocation = { level: 0, index: 0 };
     // If there's only 1 block, its block root proof will be stored at the root.
-    return this.totalNumBlocks === 1
-      ? [this.blockRootOrMergeProvingOutputs.getNode(rootLocation), this.paddingBlockRootProvingOutput]
-      : this.blockRootOrMergeProvingOutputs.getChildren(rootLocation);
-  }
-
-  #getPreviousRollupData({
-    inputs,
-    proof,
-    verificationKey,
-  }: PublicInputsAndRecursiveProof<
-    BlockRootOrBlockMergePublicInputs,
-    typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
-  >) {
-    const leafIndex = getVKIndex(verificationKey.keyAsFields);
-    return new PreviousRollupBlockData(
-      inputs,
-      proof,
-      verificationKey.keyAsFields,
-      new MembershipWitness(VK_TREE_HEIGHT, BigInt(leafIndex), getVKSiblingPath(leafIndex)),
-    );
+    return this.totalNumCheckpoints === 1
+      ? [this.checkpointProofs.getNode(rootLocation)?.provingOutput, this.checkpointPaddingProof?.provingOutput]
+      : this.checkpointProofs.getChildren(rootLocation).map(c => c?.provingOutput);
   }
 }