@aztec/prover-client 0.55.1 → 0.56.0

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

@@ -0,0 +1,232 @@
+import { type MerkleTreeId, type ProvingResult } from '@aztec/circuit-types';
+import {
+  type ARCHIVE_HEIGHT,
+  type AppendOnlyTreeSnapshot,
+  type BlockRootOrBlockMergePublicInputs,
+  Fr,
+  type GlobalVariables,
+  type L1_TO_L2_MSG_SUBTREE_SIBLING_PATH_LENGTH,
+  type NESTED_RECURSIVE_PROOF_LENGTH,
+  NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
+  type Proof,
+  type RecursiveProof,
+  type RootRollupPublicInputs,
+  type VerificationKeyAsFields,
+} from '@aztec/circuits.js';
+import { padArrayEnd } from '@aztec/foundation/collection';
+import { type Tuple } from '@aztec/foundation/serialize';
+
+import { BlockProvingState } from './block-proving-state.js';
+
+export type TreeSnapshots = Map<MerkleTreeId, AppendOnlyTreeSnapshot>;
+
+enum PROVING_STATE_LIFECYCLE {
+  PROVING_STATE_CREATED,
+  PROVING_STATE_FULL,
+  PROVING_STATE_RESOLVED,
+  PROVING_STATE_REJECTED,
+}
+
+export type BlockMergeRollupInputData = {
+  inputs: [BlockRootOrBlockMergePublicInputs | undefined, BlockRootOrBlockMergePublicInputs | undefined];
+  proofs: [
+    RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined,
+    RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined,
+  ];
+  verificationKeys: [VerificationKeyAsFields | undefined, VerificationKeyAsFields | undefined];
+};
+
+/**
+ * The current state of the proving schedule for an epoch.
+ * Contains the raw inputs and intermediate state to generate every constituent proof in the tree.
+ * Carries an identifier so we can identify if the proving state is discarded and a new one started.
+ * Captures resolve and reject callbacks to provide a promise base interface to the consumer of our proving.
+ */
+export class EpochProvingState {
+  private provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED;
+
+  private mergeRollupInputs: BlockMergeRollupInputData[] = [];
+  public rootRollupPublicInputs: RootRollupPublicInputs | undefined;
+  public finalProof: Proof | undefined;
+  public blocks: BlockProvingState[] = [];
+
+  constructor(
+    public readonly epochNumber: 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[this.blocks.length - 1];
+  }
+
+  // Returns the number of levels of merge rollups
+  public get numMergeLevels() {
+    return BigInt(Math.ceil(Math.log2(this.totalNumBlocks)) - 1);
+  }
+
+  // Calculates the index and level of the parent rollup circuit
+  // Based on tree implementation in unbalanced_tree.ts -> batchInsert()
+  // REFACTOR: This is repeated from the block orchestrator
+  public findMergeLevel(currentLevel: bigint, currentIndex: bigint) {
+    const moveUpMergeLevel = (levelSize: number, index: bigint, nodeToShift: boolean) => {
+      levelSize /= 2;
+      if (levelSize & 1) {
+        [levelSize, nodeToShift] = nodeToShift ? [levelSize + 1, false] : [levelSize - 1, true];
+      }
+      index >>= 1n;
+      return { thisLevelSize: levelSize, thisIndex: index, shiftUp: nodeToShift };
+    };
+    let [thisLevelSize, shiftUp] =
+      this.totalNumBlocks & 1 ? [this.totalNumBlocks - 1, true] : [this.totalNumBlocks, false];
+    const maxLevel = this.numMergeLevels + 1n;
+    let placeholder = currentIndex;
+    for (let i = 0; i < maxLevel - currentLevel; i++) {
+      ({ thisLevelSize, thisIndex: placeholder, shiftUp } = moveUpMergeLevel(thisLevelSize, placeholder, shiftUp));
+    }
+    let thisIndex = currentIndex;
+    let mergeLevel = currentLevel;
+    while (thisIndex >= thisLevelSize && mergeLevel != 0n) {
+      mergeLevel -= 1n;
+      ({ thisLevelSize, thisIndex, shiftUp } = moveUpMergeLevel(thisLevelSize, thisIndex, shiftUp));
+    }
+    return [mergeLevel - 1n, thisIndex >> 1n, thisIndex & 1n];
+  }
+
+  // Adds a block to the proving state, returns its index
+  // Will update the proving life cycle if this is the last block
+  public startNewBlock(
+    numTxs: number,
+    globalVariables: GlobalVariables,
+    l1ToL2Messages: Fr[],
+    messageTreeSnapshot: AppendOnlyTreeSnapshot,
+    messageTreeRootSiblingPath: Tuple<Fr, typeof L1_TO_L2_MSG_SUBTREE_SIBLING_PATH_LENGTH>,
+    messageTreeSnapshotAfterInsertion: AppendOnlyTreeSnapshot,
+    archiveTreeSnapshot: AppendOnlyTreeSnapshot,
+    archiveTreeRootSiblingPath: Tuple<Fr, typeof ARCHIVE_HEIGHT>,
+    previousBlockHash: Fr,
+    completionCallback?: (result: ProvingResult) => void,
+    rejectionCallback?: (reason: string) => void,
+  ) {
+    const block = new BlockProvingState(
+      this.blocks.length,
+      numTxs,
+      globalVariables,
+      padArrayEnd(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP),
+      messageTreeSnapshot,
+      messageTreeRootSiblingPath,
+      messageTreeSnapshotAfterInsertion,
+      archiveTreeSnapshot,
+      archiveTreeRootSiblingPath,
+      previousBlockHash,
+      completionCallback,
+      reason => {
+        // Reject the block
+        if (rejectionCallback) {
+          rejectionCallback(reason);
+        }
+        // An error on any block rejects this whole epoch
+        this.reject(reason);
+      },
+    );
+    this.blocks.push(block);
+    if (this.blocks.length === this.totalNumBlocks) {
+      this.provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_FULL;
+    }
+    return this.blocks.length - 1;
+  }
+
+  // Returns true if this proving state is still valid, false otherwise
+  public verifyState() {
+    return (
+      this.provingStateLifecycle === PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED ||
+      this.provingStateLifecycle === PROVING_STATE_LIFECYCLE.PROVING_STATE_FULL
+    );
+  }
+
+  // Returns true if we are still able to accept blocks, false otherwise
+  public isAcceptingBlocks() {
+    return this.provingStateLifecycle === PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED;
+  }
+
+  /**
+   * Stores the inputs to a merge circuit and determines if the circuit is ready to be executed
+   * @param mergeInputs - The inputs to store
+   * @param indexWithinMerge - The index in the set of inputs to this merge circuit
+   * @param indexOfMerge - The global index of this merge circuit
+   * @returns True if the merge circuit is ready to be executed, false otherwise
+   */
+  public storeMergeInputs(
+    mergeInputs: [
+      BlockRootOrBlockMergePublicInputs,
+      RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH>,
+      VerificationKeyAsFields,
+    ],
+    indexWithinMerge: number,
+    indexOfMerge: number,
+  ) {
+    if (!this.mergeRollupInputs[indexOfMerge]) {
+      const mergeInputData: BlockMergeRollupInputData = {
+        inputs: [undefined, undefined],
+        proofs: [undefined, undefined],
+        verificationKeys: [undefined, undefined],
+      };
+      mergeInputData.inputs[indexWithinMerge] = mergeInputs[0];
+      mergeInputData.proofs[indexWithinMerge] = mergeInputs[1];
+      mergeInputData.verificationKeys[indexWithinMerge] = mergeInputs[2];
+      this.mergeRollupInputs[indexOfMerge] = mergeInputData;
+      return false;
+    }
+    const mergeInputData = this.mergeRollupInputs[indexOfMerge];
+    mergeInputData.inputs[indexWithinMerge] = mergeInputs[0];
+    mergeInputData.proofs[indexWithinMerge] = mergeInputs[1];
+    mergeInputData.verificationKeys[indexWithinMerge] = mergeInputs[2];
+    return true;
+  }
+
+  // Returns a specific transaction proving state
+  public getBlockProvingState(index: number) {
+    return this.blocks[index];
+  }
+
+  // Returns a set of merge rollup inputs
+  public getMergeInputs(indexOfMerge: number) {
+    return this.mergeRollupInputs[indexOfMerge];
+  }
+
+  // Returns true if we have sufficient inputs to execute the block root rollup
+  public isReadyForRootRollup() {
+    return !(this.mergeRollupInputs[0] === undefined || this.mergeRollupInputs[0].inputs.findIndex(p => !p) !== -1);
+  }
+
+  // Attempts to reject the proving state promise with a reason of 'cancelled'
+  public cancel() {
+    this.reject('Proving cancelled');
+  }
+
+  // Attempts to reject the proving state promise with the given reason
+  // Does nothing if not in a valid state
+  public reject(reason: string) {
+    if (!this.verifyState()) {
+      return;
+    }
+    this.provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_REJECTED;
+    this.rejectionCallback(reason);
+
+    for (const block of this.blocks) {
+      block.reject('Proving cancelled');
+    }
+  }
+
+  // Attempts to resolve the proving state promise with the given result
+  // Does nothing if not in a valid state
+  public resolve(result: ProvingResult) {
+    if (!this.verifyState()) {
+      return;
+    }
+    this.provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_RESOLVED;
+    this.completionCallback(result);
+  }
+}