@aztec/prover-client 0.42.0

package/src/orchestrator/proving-state.ts

@@ -0,0 +1,210 @@
+import { type L2Block, type MerkleTreeId, type ProvingResult } from '@aztec/circuit-types';
+import {
+  type AppendOnlyTreeSnapshot,
+  type BaseOrMergeRollupPublicInputs,
+  type Fr,
+  type GlobalVariables,
+  type L1_TO_L2_MSG_SUBTREE_SIBLING_PATH_LENGTH,
+  type NESTED_RECURSIVE_PROOF_LENGTH,
+  type NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
+  type Proof,
+  type RECURSIVE_PROOF_LENGTH,
+  type RecursiveProof,
+  type RootParityInput,
+  type RootRollupPublicInputs,
+  type VerificationKeyAsFields,
+  type VerificationKeys,
+} from '@aztec/circuits.js';
+import { type Tuple } from '@aztec/foundation/serialize';
+
+import { type TxProvingState } from './tx-proving-state.js';
+
+export type MergeRollupInputData = {
+  inputs: [BaseOrMergeRollupPublicInputs | undefined, BaseOrMergeRollupPublicInputs | undefined];
+  proofs: [
+    RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined,
+    RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined,
+  ];
+  verificationKeys: [VerificationKeyAsFields | undefined, VerificationKeyAsFields | undefined];
+};
+
+export type TreeSnapshots = Map<MerkleTreeId, AppendOnlyTreeSnapshot>;
+
+enum PROVING_STATE_LIFECYCLE {
+  PROVING_STATE_CREATED,
+  PROVING_STATE_FULL,
+  PROVING_STATE_RESOLVED,
+  PROVING_STATE_REJECTED,
+}
+
+/**
+ * The current state of the proving schedule. Contains the raw inputs (txs) 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 ProvingState {
+  private provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED;
+  private mergeRollupInputs: MergeRollupInputData[] = [];
+  private rootParityInputs: Array<RootParityInput<typeof RECURSIVE_PROOF_LENGTH> | undefined> = [];
+  private finalRootParityInputs: RootParityInput<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined;
+  public rootRollupPublicInputs: RootRollupPublicInputs | undefined;
+  public finalAggregationObject: Fr[] | undefined;
+  public finalProof: Proof | undefined;
+  public block: L2Block | undefined;
+  private txs: TxProvingState[] = [];
+  constructor(
+    public readonly totalNumTxs: number,
+    private completionCallback: (result: ProvingResult) => void,
+    private rejectionCallback: (reason: string) => void,
+    public readonly globalVariables: GlobalVariables,
+    public readonly newL1ToL2Messages: Tuple<Fr, typeof NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP>,
+    numRootParityInputs: number,
+    public readonly messageTreeSnapshot: AppendOnlyTreeSnapshot,
+    public readonly messageTreeRootSiblingPath: Tuple<Fr, typeof L1_TO_L2_MSG_SUBTREE_SIBLING_PATH_LENGTH>,
+    public readonly privateKernelVerificationKeys: VerificationKeys,
+  ) {
+    this.rootParityInputs = Array.from({ length: numRootParityInputs }).map(_ => undefined);
+  }
+
+  // Returns the number of levels of merge rollups
+  public get numMergeLevels() {
+    return BigInt(Math.ceil(Math.log2(this.totalNumTxs)) - 1);
+  }
+
+  // Adds a transaction to the proving state, returns it's index
+  // Will update the proving life cycle if this is the last transaction
+  public addNewTx(tx: TxProvingState) {
+    this.txs.push(tx);
+    if (this.txs.length === this.totalNumTxs) {
+      this.provingStateLifecycle = PROVING_STATE_LIFECYCLE.PROVING_STATE_FULL;
+    }
+    return this.txs.length - 1;
+  }
+
+  // Returns the number of received transactions
+  public get transactionsReceived() {
+    return this.txs.length;
+  }
+
+  // Returns the final set of root parity inputs
+  public get finalRootParityInput() {
+    return this.finalRootParityInputs;
+  }
+
+  // Sets the final set of root parity inputs
+  public set finalRootParityInput(input: RootParityInput<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined) {
+    this.finalRootParityInputs = input;
+  }
+
+  // Returns the set of root parity inputs
+  public get rootParityInput() {
+    return this.rootParityInputs;
+  }
+
+  // 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 transactions, false otherwise
+  public isAcceptingTransactions() {
+    return this.provingStateLifecycle === PROVING_STATE_LIFECYCLE.PROVING_STATE_CREATED;
+  }
+
+  // Returns the complete set of transaction proving state objects
+  public get allTxs() {
+    return this.txs;
+  }
+
+  /**
+   * 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: [
+      BaseOrMergeRollupPublicInputs,
+      RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH>,
+      VerificationKeyAsFields,
+    ],
+    indexWithinMerge: number,
+    indexOfMerge: number,
+  ) {
+    if (!this.mergeRollupInputs[indexOfMerge]) {
+      const mergeInputData: MergeRollupInputData = {
+        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 getTxProvingState(txIndex: number) {
+    return this.txs[txIndex];
+  }
+
+  // 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 root rollup
+  public isReadyForRootRollup() {
+    return !(
+      this.mergeRollupInputs[0] === undefined ||
+      this.finalRootParityInput === undefined ||
+      this.mergeRollupInputs[0].inputs.findIndex(p => !p) !== -1
+    );
+  }
+
+  // Stores a set of root parity inputs at the given index
+  public setRootParityInputs(inputs: RootParityInput<typeof RECURSIVE_PROOF_LENGTH>, index: number) {
+    this.rootParityInputs[index] = inputs;
+  }
+
+  // Returns true if we have sufficient root parity inputs to execute the root parity circuit
+  public areRootParityInputsReady() {
+    return this.rootParityInputs.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);
+  }
+
+  // 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);
+  }
+}

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

@@ -0,0 +1,139 @@
+import {
+  AVM_REQUEST,
+  type AvmProvingRequest,
+  type MerkleTreeId,
+  type ProcessedTx,
+  type PublicKernelRequest,
+  PublicKernelType,
+} from '@aztec/circuit-types';
+import {
+  type AppendOnlyTreeSnapshot,
+  type BaseRollupInputs,
+  NESTED_RECURSIVE_PROOF_LENGTH,
+  type Proof,
+  type RecursiveProof,
+  type VerificationKeyData,
+  makeRecursiveProofFromBinary,
+} from '@aztec/circuits.js';
+
+export enum TX_PROVING_CODE {
+  NOT_READY,
+  READY,
+  COMPLETED,
+}
+
+export type PublicFunction = {
+  vmRequest: AvmProvingRequest | undefined;
+  vmProof: Proof | undefined;
+  previousProofType: PublicKernelType;
+  previousKernelProven: boolean;
+  publicKernelRequest: PublicKernelRequest;
+};
+
+// Type encapsulating the instruction to the orchestrator as to what
+// needs to be proven next
+export type TxProvingInstruction = {
+  code: TX_PROVING_CODE;
+  function: PublicFunction | undefined;
+};
+
+/**
+ * Helper class to manage the proving cycle of a transaction
+ * This includes the public VMs and the public kernels
+ * Also stores the inputs to the base rollup for this transaction and the tree snapshots
+ */
+export class TxProvingState {
+  private publicFunctions: PublicFunction[] = [];
+
+  constructor(
+    public readonly processedTx: ProcessedTx,
+    public readonly baseRollupInputs: BaseRollupInputs,
+    public readonly treeSnapshots: Map<MerkleTreeId, AppendOnlyTreeSnapshot>,
+    privateKernelVk: VerificationKeyData,
+  ) {
+    let previousKernelProof: RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined =
+      makeRecursiveProofFromBinary(processedTx.proof, NESTED_RECURSIVE_PROOF_LENGTH);
+    let previousProofType = PublicKernelType.NON_PUBLIC;
+    for (let i = 0; i < processedTx.publicProvingRequests.length; i++) {
+      const provingRequest = processedTx.publicProvingRequests[i];
+      const kernelRequest = provingRequest.type === AVM_REQUEST ? provingRequest.kernelRequest : provingRequest;
+      // the first circuit has a valid previous proof, it came from private
+      if (previousKernelProof) {
+        kernelRequest.inputs.previousKernel.proof = previousKernelProof;
+        kernelRequest.inputs.previousKernel.vk = privateKernelVk;
+      }
+      const vmRequest = provingRequest.type === AVM_REQUEST ? provingRequest : undefined;
+      const publicFunction: PublicFunction = {
+        vmRequest,
+        vmProof: undefined,
+        previousProofType,
+        previousKernelProven: i === 0,
+        publicKernelRequest: kernelRequest,
+      };
+      this.publicFunctions.push(publicFunction);
+      previousKernelProof = undefined;
+      previousProofType = kernelRequest.type;
+    }
+  }
+
+  // Updates the transaction's proving state after completion of a kernel proof
+  // Returns an instruction as to the next stage of tx proving
+  public getNextPublicKernelFromKernelProof(
+    provenIndex: number,
+    proof: RecursiveProof<typeof NESTED_RECURSIVE_PROOF_LENGTH>,
+    verificationKey: VerificationKeyData,
+  ): TxProvingInstruction {
+    const kernelRequest = this.getPublicFunctionState(provenIndex).publicKernelRequest;
+    const nextKernelIndex = provenIndex + 1;
+    if (nextKernelIndex >= this.publicFunctions.length) {
+      // The next kernel index is greater than our set of functions, we are done!
+      return { code: TX_PROVING_CODE.COMPLETED, function: undefined };
+    }
+
+    // There is more work to do, are we ready?
+    const nextFunction = this.publicFunctions[nextKernelIndex];
+
+    // pass both the proof and verification key forward to the next circuit
+    nextFunction.publicKernelRequest.inputs.previousKernel.proof = proof;
+    nextFunction.publicKernelRequest.inputs.previousKernel.vk = verificationKey;
+
+    // We need to update this so the state machine knows this proof is ready
+    nextFunction.previousKernelProven = true;
+    nextFunction.previousProofType = kernelRequest.type;
+    if (nextFunction.vmProof === undefined) {
+      // The VM proof for the next function is not ready
+      return { code: TX_PROVING_CODE.NOT_READY, function: undefined };
+    }
+
+    // The VM proof is ready, we can continue
+    return { code: TX_PROVING_CODE.READY, function: nextFunction };
+  }
+
+  // Updates the transaction's proving state after completion of a VM proof
+  // Returns an instruction as to the next stage of tx proving
+  public getNextPublicKernelFromVMProof(provenIndex: number, proof: Proof): TxProvingInstruction {
+    const provenFunction = this.publicFunctions[provenIndex];
+    provenFunction.vmProof = proof;
+
+    if (!provenFunction.previousKernelProven) {
+      // The previous kernel is not yet ready
+      return { code: TX_PROVING_CODE.NOT_READY, function: undefined };
+    }
+    // The previous kernel is ready so we can prove this kernel
+    return { code: TX_PROVING_CODE.READY, function: provenFunction };
+  }
+
+  // Returns the public function state at the given index
+  // Throws if out of bounds
+  public getPublicFunctionState(functionIndex: number) {
+    if (functionIndex < 0 || functionIndex >= this.publicFunctions.length) {
+      throw new Error(`Requested public function index was out of bounds`);
+    }
+    return this.publicFunctions[functionIndex];
+  }
+
+  // Returns the number of public kernels required by this transaction
+  public getNumPublicKernels() {
+    return this.publicFunctions.length;
+  }
+}

package/src/prover-agent/index.ts

@@ -0,0 +1,3 @@
+export * from './prover-agent.js';
+export * from './memory-proving-queue.js';
+export * from './rpc.js';

package/src/prover-agent/memory-proving-queue.ts

@@ -0,0 +1,303 @@
+import {
+  type ProofAndVerificationKey,
+  type ProvingJob,
+  type ProvingJobSource,
+  type ProvingRequest,
+  type ProvingRequestResult,
+  ProvingRequestType,
+  type PublicInputsAndRecursiveProof,
+  type PublicKernelNonTailRequest,
+  type PublicKernelTailRequest,
+  type ServerCircuitProver,
+} from '@aztec/circuit-types';
+import type {
+  AvmCircuitInputs,
+  BaseOrMergeRollupPublicInputs,
+  BaseParityInputs,
+  BaseRollupInputs,
+  KernelCircuitPublicInputs,
+  MergeRollupInputs,
+  NESTED_RECURSIVE_PROOF_LENGTH,
+  PrivateKernelEmptyInputData,
+  PublicKernelCircuitPublicInputs,
+  RECURSIVE_PROOF_LENGTH,
+  RootParityInput,
+  RootParityInputs,
+  RootRollupInputs,
+  RootRollupPublicInputs,
+} from '@aztec/circuits.js';
+import { randomBytes } from '@aztec/foundation/crypto';
+import { AbortedError, TimeoutError } from '@aztec/foundation/error';
+import { MemoryFifo } from '@aztec/foundation/fifo';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { type PromiseWithResolvers, promiseWithResolvers } from '@aztec/foundation/promise';
+
+type ProvingJobWithResolvers<T extends ProvingRequest = ProvingRequest> = {
+  id: string;
+  request: T;
+  signal?: AbortSignal;
+  attempts: number;
+} & PromiseWithResolvers<ProvingRequestResult<T['type']>>;
+
+const MAX_RETRIES = 3;
+
+const defaultIdGenerator = () => randomBytes(4).toString('hex');
+
+/**
+ * A helper class that sits in between services that need proofs created and agents that can create them.
+ * The queue accumulates jobs and provides them to agents in FIFO order.
+ */
+export class MemoryProvingQueue implements ServerCircuitProver, ProvingJobSource {
+  private log = createDebugLogger('aztec:prover-client:prover-pool:queue');
+  private queue = new MemoryFifo<ProvingJobWithResolvers>();
+  private jobsInProgress = new Map<string, ProvingJobWithResolvers>();
+
+  constructor(private generateId = defaultIdGenerator) {}
+
+  async getProvingJob({ timeoutSec = 1 } = {}): Promise<ProvingJob<ProvingRequest> | undefined> {
+    try {
+      const job = await this.queue.get(timeoutSec);
+      if (!job) {
+        return undefined;
+      }
+
+      if (job.signal?.aborted) {
+        this.log.debug(`Job ${job.id} type=${job.request.type} has been aborted`);
+        return undefined;
+      }
+
+      this.jobsInProgress.set(job.id, job);
+      return {
+        id: job.id,
+        request: job.request,
+      };
+    } catch (err) {
+      if (err instanceof TimeoutError) {
+        return undefined;
+      }
+
+      throw err;
+    }
+  }
+
+  resolveProvingJob<T extends ProvingRequestType>(jobId: string, result: ProvingRequestResult<T>): Promise<void> {
+    const job = this.jobsInProgress.get(jobId);
+    if (!job) {
+      return Promise.reject(new Error('Job not found'));
+    }
+
+    this.jobsInProgress.delete(jobId);
+
+    if (job.signal?.aborted) {
+      return Promise.resolve();
+    }
+
+    job.resolve(result);
+    return Promise.resolve();
+  }
+
+  rejectProvingJob(jobId: string, err: any): Promise<void> {
+    const job = this.jobsInProgress.get(jobId);
+    if (!job) {
+      return Promise.reject(new Error('Job not found'));
+    }
+
+    this.jobsInProgress.delete(jobId);
+
+    if (job.signal?.aborted) {
+      return Promise.resolve();
+    }
+
+    if (job.attempts < MAX_RETRIES) {
+      job.attempts++;
+      this.log.warn(
+        `Job id=${job.id} type=${ProvingRequestType[job.request.type]} failed with error: ${err}. Retry ${
+          job.attempts
+        }/${MAX_RETRIES}`,
+      );
+      this.queue.put(job);
+    } else {
+      this.log.error(`Job id=${job.id} type=${ProvingRequestType[job.request.type]} failed with error: ${err}`);
+      job.reject(err);
+    }
+    return Promise.resolve();
+  }
+
+  private enqueue<T extends ProvingRequest>(
+    request: T,
+    signal?: AbortSignal,
+  ): Promise<ProvingRequestResult<T['type']>> {
+    const { promise, resolve, reject } = promiseWithResolvers<ProvingRequestResult<T['type']>>();
+    const item: ProvingJobWithResolvers<T> = {
+      id: this.generateId(),
+      request,
+      signal,
+      promise,
+      resolve,
+      reject,
+      attempts: 1,
+    };
+
+    if (signal) {
+      signal.addEventListener('abort', () => reject(new AbortedError('Operation has been aborted')));
+    }
+
+    this.log.debug(
+      `Adding id=${item.id} type=${ProvingRequestType[request.type]} proving job to queue depth=${this.queue.length()}`,
+    );
+    // TODO (alexg) remove the `any`
+    if (!this.queue.put(item as any)) {
+      throw new Error();
+    }
+
+    return promise;
+  }
+
+  getEmptyPrivateKernelProof(
+    inputs: PrivateKernelEmptyInputData,
+    signal?: AbortSignal,
+  ): Promise<PublicInputsAndRecursiveProof<KernelCircuitPublicInputs>> {
+    return this.enqueue({ type: ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs }, signal);
+  }
+
+  /**
+   * Creates a proof for the given input.
+   * @param input - Input to the circuit.
+   */
+  getBaseParityProof(
+    inputs: BaseParityInputs,
+    signal?: AbortSignal,
+  ): Promise<RootParityInput<typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.BASE_PARITY,
+        inputs,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Creates a proof for the given input.
+   * @param input - Input to the circuit.
+   */
+  getRootParityProof(
+    inputs: RootParityInputs,
+    signal?: AbortSignal,
+  ): Promise<RootParityInput<typeof NESTED_RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.ROOT_PARITY,
+        inputs,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Creates a proof for the given input.
+   * @param input - Input to the circuit.
+   */
+  getBaseRollupProof(
+    input: BaseRollupInputs,
+    signal?: AbortSignal,
+  ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.BASE_ROLLUP,
+        inputs: input,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Creates a proof for the given input.
+   * @param input - Input to the circuit.
+   */
+  getMergeRollupProof(
+    input: MergeRollupInputs,
+    signal?: AbortSignal,
+  ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.MERGE_ROLLUP,
+        inputs: input,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Creates a proof for the given input.
+   * @param input - Input to the circuit.
+   */
+  getRootRollupProof(
+    input: RootRollupInputs,
+    signal?: AbortSignal,
+  ): Promise<PublicInputsAndRecursiveProof<RootRollupPublicInputs>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.ROOT_ROLLUP,
+        inputs: input,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Create a public kernel proof.
+   * @param kernelRequest - Object containing the details of the proof required
+   */
+  getPublicKernelProof(
+    kernelRequest: PublicKernelNonTailRequest,
+    signal?: AbortSignal,
+  ): Promise<PublicInputsAndRecursiveProof<PublicKernelCircuitPublicInputs>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.PUBLIC_KERNEL_NON_TAIL,
+        kernelType: kernelRequest.type,
+        inputs: kernelRequest.inputs,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Create a public kernel tail proof.
+   * @param kernelRequest - Object containing the details of the proof required
+   */
+  getPublicTailProof(
+    kernelRequest: PublicKernelTailRequest,
+    signal?: AbortSignal,
+  ): Promise<PublicInputsAndRecursiveProof<KernelCircuitPublicInputs>> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.PUBLIC_KERNEL_TAIL,
+        kernelType: kernelRequest.type,
+        inputs: kernelRequest.inputs,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Creates an AVM proof.
+   */
+  getAvmProof(inputs: AvmCircuitInputs, signal?: AbortSignal | undefined): Promise<ProofAndVerificationKey> {
+    return this.enqueue(
+      {
+        type: ProvingRequestType.PUBLIC_VM,
+        inputs,
+      },
+      signal,
+    );
+  }
+
+  /**
+   * Verifies a circuit proof
+   */
+  verifyProof(): Promise<void> {
+    return Promise.reject('not implemented');
+  }
+}