@aztec/prover-client 0.0.1-commit.a89ec08 → 0.0.1-commit.aa0c64f

package/src/orchestrator/top-tree-proving-state.ts

@@ -0,0 +1,219 @@
+import type { BatchedBlob, BatchedBlobAccumulator, FinalBlobBatchingChallenges } from '@aztec/blob-lib';
+import type { NESTED_RECURSIVE_PROOF_LENGTH, NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH } from '@aztec/constants';
+import { EpochNumber } from '@aztec/foundation/branded-types';
+import { type TreeNodeLocation, UnbalancedTreeStore } from '@aztec/foundation/trees';
+import type { PublicInputsAndRecursiveProof } from '@aztec/stdlib/interfaces/server';
+import type { Proof } from '@aztec/stdlib/proofs';
+import {
+  CheckpointMergeRollupPrivateInputs,
+  CheckpointPaddingRollupPrivateInputs,
+  CheckpointRollupPublicInputs,
+  RootRollupPrivateInputs,
+  type RootRollupPublicInputs,
+} from '@aztec/stdlib/rollup';
+
+import { toProofData } from './block-building-helpers.js';
+import type { ProofState } from './block-proving-state.js';
+
+enum TOP_TREE_LIFECYCLE {
+  CREATED,
+  RESOLVED,
+  REJECTED,
+}
+
+/**
+ * Lean top-tree-only state. Owns the merge tree of checkpoint root proofs, the
+ * single-checkpoint padding proof slot, the final root rollup proof, and the blob
+ * accumulator endpoints needed to finalise the epoch's batched blob proof.
+ *
+ * Constructed with `totalNumCheckpoints` and `finalBlobBatchingChallenges` upfront —
+ * by the time the top tree starts, all checkpoints are known and the challenges are
+ * derivable from their blob fields.
+ */
+export class TopTreeProvingState {
+  private checkpointProofs: UnbalancedTreeStore<
+    ProofState<CheckpointRollupPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+  >;
+  private checkpointPaddingProof:
+    | ProofState<CheckpointRollupPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+    | undefined;
+  private rootRollupProof: ProofState<RootRollupPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined;
+  private endBlobAccumulator: BatchedBlobAccumulator | undefined;
+  private finalBatchedBlob: BatchedBlob | undefined;
+  private lifecycle = TOP_TREE_LIFECYCLE.CREATED;
+  constructor(
+    public readonly epochNumber: EpochNumber,
+    public readonly totalNumCheckpoints: number,
+    public readonly finalBlobBatchingChallenges: FinalBlobBatchingChallenges,
+    public readonly startBlobAccumulator: BatchedBlobAccumulator,
+    private readonly completionCallback: () => void,
+    private readonly rejectionCallback: (reason: string) => void,
+  ) {
+    if (totalNumCheckpoints < 1) {
+      throw new Error(`TopTreeProvingState requires at least one checkpoint; got ${totalNumCheckpoints}.`);
+    }
+    this.checkpointProofs = new UnbalancedTreeStore(totalNumCheckpoints);
+  }
+
+  // --- checkpoint root rollup ---
+
+  public setCheckpointRootRollupProof(
+    checkpointIndex: number,
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
+      typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
+    >,
+  ): TreeNodeLocation {
+    return this.checkpointProofs.setLeaf(checkpointIndex, { provingOutput });
+  }
+
+  // --- checkpoint merge rollup ---
+
+  public tryStartProvingCheckpointMerge(location: TreeNodeLocation) {
+    if (this.checkpointProofs.getNode(location)?.isProving) {
+      return false;
+    }
+    this.checkpointProofs.setNode(location, { isProving: true });
+    return true;
+  }
+
+  public setCheckpointMergeRollupProof(
+    location: TreeNodeLocation,
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
+      typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
+    >,
+  ) {
+    this.checkpointProofs.setNode(location, { provingOutput });
+  }
+
+  public isReadyForCheckpointMerge(location: TreeNodeLocation) {
+    return !!this.checkpointProofs.getSibling(location)?.provingOutput;
+  }
+
+  public getParentLocation(location: TreeNodeLocation) {
+    return this.checkpointProofs.getParentLocation(location);
+  }
+
+  public getCheckpointMergeRollupInputs(mergeLocation: TreeNodeLocation) {
+    const [left, right] = this.checkpointProofs.getChildren(mergeLocation).map(c => c?.provingOutput);
+    if (!left || !right) {
+      throw new Error('At least one child is not ready for the checkpoint merge rollup.');
+    }
+    return new CheckpointMergeRollupPrivateInputs([toProofData(left), toProofData(right)]);
+  }
+
+  // --- padding (single-checkpoint case) ---
+
+  public tryStartProvingPaddingCheckpoint() {
+    if (this.checkpointPaddingProof?.isProving) {
+      return false;
+    }
+    this.checkpointPaddingProof = { isProving: true };
+    return true;
+  }
+
+  public setCheckpointPaddingProof(
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
+      typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
+    >,
+  ) {
+    this.checkpointPaddingProof = { provingOutput };
+  }
+
+  public getPaddingCheckpointInputs() {
+    return new CheckpointPaddingRollupPrivateInputs();
+  }
+
+  // --- root rollup ---
+
+  public tryStartProvingRootRollup() {
+    if (this.rootRollupProof?.isProving) {
+      return false;
+    }
+    this.rootRollupProof = { isProving: true };
+    return true;
+  }
+
+  public setRootRollupProof(provingOutput: PublicInputsAndRecursiveProof<RootRollupPublicInputs>) {
+    this.rootRollupProof = { provingOutput };
+  }
+
+  public isReadyForRootRollup() {
+    const childProofs = this.#getChildProofsForRoot();
+    return childProofs.every(p => !!p);
+  }
+
+  public getRootRollupInputs() {
+    const [left, right] = this.#getChildProofsForRoot();
+    if (!left || !right) {
+      throw new Error('At least one child is not ready for the root rollup.');
+    }
+    return RootRollupPrivateInputs.from({
+      previousRollups: [toProofData(left), toProofData(right)],
+    });
+  }
+
+  // --- blob accumulator finalisation ---
+
+  /**
+   * Sets the end-of-epoch blob accumulator, computed by the top-tree orchestrator
+   * from the surviving checkpoints' blob fields. Required before `finalizeBatchedBlob`.
+   */
+  public setEndBlobAccumulator(accumulator: BatchedBlobAccumulator) {
+    this.endBlobAccumulator = accumulator;
+  }
+
+  public async finalizeBatchedBlob() {
+    if (!this.endBlobAccumulator) {
+      throw new Error('End blob accumulator not set; call setEndBlobAccumulator before finalize.');
+    }
+    this.finalBatchedBlob = await this.endBlobAccumulator.finalize(true /* verifyProof */);
+  }
+
+  public getEpochProofResult(): { proof: Proof; publicInputs: RootRollupPublicInputs; batchedBlobInputs: BatchedBlob } {
+    const provingOutput = this.rootRollupProof?.provingOutput;
+    if (!provingOutput || !this.finalBatchedBlob) {
+      throw new Error('Top-tree proof not ready; root rollup or batched blob missing.');
+    }
+    return {
+      proof: provingOutput.proof.binaryProof,
+      publicInputs: provingOutput.inputs,
+      batchedBlobInputs: this.finalBatchedBlob,
+    };
+  }
+
+  // --- lifecycle ---
+
+  public verifyState() {
+    return this.lifecycle === TOP_TREE_LIFECYCLE.CREATED;
+  }
+
+  public resolve() {
+    if (!this.verifyState()) {
+      return;
+    }
+    this.lifecycle = TOP_TREE_LIFECYCLE.RESOLVED;
+    this.completionCallback();
+  }
+
+  public reject(reason: string) {
+    if (!this.verifyState()) {
+      return;
+    }
+    this.lifecycle = TOP_TREE_LIFECYCLE.REJECTED;
+    this.rejectionCallback(reason);
+  }
+
+  public cancel() {
+    this.reject('Top-tree proving cancelled');
+  }
+
+  #getChildProofsForRoot() {
+    const rootLocation = { level: 0, index: 0 };
+    return this.totalNumCheckpoints === 1
+      ? [this.checkpointProofs.getNode(rootLocation)?.provingOutput, this.checkpointPaddingProof?.provingOutput]
+      : this.checkpointProofs.getChildren(rootLocation).map(c => c?.provingOutput);
+  }
+}

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

@@ -1,4 +1,4 @@
-import { AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED, NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH } from '@aztec/constants';
+import { AVM_V2_PROOF_LENGTH_IN_FIELDS, NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH } from '@aztec/constants';
 import type { Fr } from '@aztec/foundation/curves/bn254';
 import { getVkData } from '@aztec/noir-protocol-circuits-types/server/vks';
 import type { AvmCircuitInputs } from '@aztec/stdlib/avm';
@@ -32,7 +32,7 @@ export class TxProvingState {
     PublicChonkVerifierPublicInputs,
     typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
   >;
-  private avmProof?: RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>;
+  private avmProof?: RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS>;
 
   constructor(
     public readonly processedTx: ProcessedTx,
@@ -80,7 +80,7 @@ export class TxProvingState {
     this.publicChonkVerifier = publicChonkVerifierProofAndVk;
   }
 
-  public setAvmProof(avmProof: RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>) {
+  public setAvmProof(avmProof: RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS>) {
     this.avmProof = avmProof;
   }
 

package/src/prover-client/prover-client.ts

@@ -1,11 +1,13 @@
 import { type ACVMConfig, type BBConfig, BBNativeRollupProver, TestCircuitProver } from '@aztec/bb-prover';
+import type { EpochNumber } from '@aztec/foundation/branded-types';
 import { times } from '@aztec/foundation/collection';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { EthAddress } from '@aztec/foundation/eth-address';
 import { type Logger, createLogger } from '@aztec/foundation/log';
+import { SerialQueue } from '@aztec/foundation/queue';
 import { NativeACVMSimulator } from '@aztec/simulator/server';
 import {
   type ActualProverConfig,
-  type EpochProver,
   type EpochProverManager,
   type ForkMerkleTreeOperations,
   type ProvingJobBroker,
@@ -15,19 +17,68 @@ import {
   type ServerCircuitProver,
   tryStop,
 } from '@aztec/stdlib/interfaces/server';
+import type { CheckpointConstantData } from '@aztec/stdlib/rollup';
+import type { BlockHeader } from '@aztec/stdlib/tx';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import type { ProverClientConfig } from '../config.js';
-import { ProvingOrchestrator } from '../orchestrator/orchestrator.js';
+import { CheckpointSubTreeOrchestrator } from '../orchestrator/checkpoint-sub-tree-orchestrator.js';
+import type { ChonkCache } from '../orchestrator/chonk-cache.js';
+import { TopTreeOrchestrator } from '../orchestrator/top-tree-orchestrator.js';
 import { BrokerCircuitProverFacade } from '../proving_broker/broker_prover_facade.js';
 import { InlineProofStore, type ProofStore, createProofStore } from '../proving_broker/proof_store/index.js';
 import { ProvingAgent } from '../proving_broker/proving_agent.js';
-import { ServerEpochProver } from './server-epoch-prover.js';
+
+/**
+ * The factory surface that `EpochProvingJob` (in `prover-node`) depends on. Implemented
+ * by `ProverClient`. Defined here rather than in stdlib because the return types
+ * (`CheckpointSubTreeOrchestrator`, `TopTreeOrchestrator`) are concrete classes from
+ * this package.
+ *
+ * A single `BrokerCircuitProverFacade` is owned by `ProverClient` and shared across
+ * every orchestrator (every sub-tree and every top-tree across every concurrent epoch
+ * job). The broker delivers each completed-job notification exactly once (drained on
+ * the first `getCompletedJobs` poll), so multiple facades polling the same broker
+ * race and lose notifications
+ *
+ * The facade's job map cleans up entries on resolve/reject, and the prover-node
+ * keeps `ProverClient` alive for its whole lifetime
+ */
+export interface EpochProverFactory {
+  getProverId(): EthAddress;
+  /**
+   * Constructs and starts a `CheckpointSubTreeOrchestrator` for a single checkpoint
+   * against the supplied shared `chonkCache`. The cache is owned by the prover-node
+   * and survives across epochs / sessions.
+   */
+  createCheckpointSubTreeOrchestrator(
+    chonkCache: ChonkCache,
+    epochNumber: EpochNumber,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator>;
+  createTopTreeOrchestrator(): TopTreeOrchestrator;
+}
 
 /** Manages proving of epochs by orchestrating the proving of individual blocks relying on a pool of prover agents. */
-export class ProverClient implements EpochProverManager {
+export class ProverClient implements EpochProverManager, EpochProverFactory {
   private running = false;
   private agents: ProvingAgent[] = [];
+  /**
+   * The single broker facade shared by every orchestrator created from this client.
+   * Constructed lazily on `start()` and torn down on `stop()` — see the comment on
+   * `EpochProverFactory` for why a single shared facade is required.
+   */
+  private facade: BrokerCircuitProverFacade | undefined;
+  /**
+   * Single deferred-proving-job queue shared across every orchestrator (sub-trees and
+   * top-trees, across every concurrent epoch session). Throttles the total rate of job
+   * submission to the broker once, rather than once per orchestrator. Started lazily
+   * alongside the facade and cancelled on `stop()`.
+   */
+  private deferredJobQueue: SerialQueue | undefined;
 
   private constructor(
     private config: ProverClientConfig,
@@ -40,25 +91,76 @@ export class ProverClient implements EpochProverManager {
     private log: Logger = createLogger('prover-client:tx-prover'),
   ) {}
 
-  public createEpochProver(): EpochProver {
-    const bindings = this.log.getBindings();
-    const facade = new BrokerCircuitProverFacade(
-      this.orchestratorClient,
-      this.proofStore,
-      this.failedProofStore,
-      undefined,
-      bindings,
-    );
-    const orchestrator = new ProvingOrchestrator(
+  /**
+   * Lazy-init the shared facade. The broker delivers each completed-job notification
+   * exactly once (drained on the first `getCompletedJobs` poll), so we cannot start
+   * a shared facade alongside the per-call facades that `createEpochProver` builds —
+   * they would race for notifications and one side would silently drop them. Starting
+   * the shared facade only on first use of one of the new factory methods keeps the
+   * legacy `createEpochProver` path race-free.
+   */
+  private getFacade(): BrokerCircuitProverFacade {
+    if (!this.running) {
+      throw new Error('ProverClient is not running; call start() before constructing orchestrators.');
+    }
+    if (!this.facade) {
+      this.facade = new BrokerCircuitProverFacade(
+        this.orchestratorClient,
+        this.proofStore,
+        this.failedProofStore,
+        undefined,
+        this.log.getBindings(),
+      );
+      this.facade.start();
+    }
+    return this.facade;
+  }
+
+  /** Lazy-init the shared deferred-job queue, started with the configured enqueue concurrency. */
+  private getDeferredJobQueue(): SerialQueue {
+    if (!this.running) {
+      throw new Error('ProverClient is not running; call start() before constructing orchestrators.');
+    }
+    if (!this.deferredJobQueue) {
+      this.deferredJobQueue = new SerialQueue();
+      this.deferredJobQueue.start(this.config.enqueueConcurrency);
+    }
+    return this.deferredJobQueue;
+  }
+
+  public createCheckpointSubTreeOrchestrator(
+    chonkCache: ChonkCache,
+    epochNumber: EpochNumber,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator> {
+    return CheckpointSubTreeOrchestrator.start(
       this.worldState,
-      facade,
+      this.getFacade(),
       this.config.proverId,
+      chonkCache,
+      epochNumber,
       this.config.cancelJobsOnStop,
-      this.config.enqueueConcurrency,
+      this.getDeferredJobQueue(),
+      checkpointConstants,
+      l1ToL2Messages,
+      totalNumBlocks,
+      headerOfLastBlockInPreviousCheckpoint,
       this.telemetry,
-      bindings,
+      this.log.getBindings(),
+    );
+  }
+
+  public createTopTreeOrchestrator(): TopTreeOrchestrator {
+    return new TopTreeOrchestrator(
+      this.getFacade(),
+      this.config.proverId,
+      this.getDeferredJobQueue(),
+      this.telemetry,
+      this.log.getBindings(),
     );
-    return new ServerEpochProver(facade, orchestrator);
   }
 
   public getProverId(): EthAddress {
@@ -100,6 +202,14 @@ export class ProverClient implements EpochProverManager {
     }
     this.running = false;
     await this.stopAgents();
+    if (this.deferredJobQueue) {
+      await this.deferredJobQueue.cancel();
+      this.deferredJobQueue = undefined;
+    }
+    if (this.facade) {
+      await tryStop(this.facade, this.log);
+      this.facade = undefined;
+    }
     await tryStop(this.orchestratorClient);
   }
 

package/src/proving_broker/broker_prover_facade.ts

@@ -1,10 +1,12 @@
 import type {
-  AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED,
+  AVM_V2_PROOF_LENGTH_IN_FIELDS,
   NESTED_RECURSIVE_PROOF_LENGTH,
   NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH,
   RECURSIVE_PROOF_LENGTH,
 } from '@aztec/constants';
+import { asyncPool } from '@aztec/foundation/async-pool';
 import { EpochNumber } from '@aztec/foundation/branded-types';
+import { chunk } from '@aztec/foundation/collection';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { truncate } from '@aztec/foundation/string';
@@ -226,17 +228,11 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
       // We collect all returned notifications and return them
       const allCompleted = new Set<ProvingJobId>();
       try {
-        let numRequests = 0;
-        while (ids.length > 0) {
-          const slice = ids.splice(0, SNAPSHOT_SYNC_CHECK_MAX_REQUEST_SIZE);
-          const completed = await this.broker.getCompletedJobs(slice);
+        const batches = ids.length > 0 ? chunk(ids, SNAPSHOT_SYNC_CHECK_MAX_REQUEST_SIZE) : [[]];
+        await asyncPool(1, batches, async batch => {
+          const completed = await this.broker.getCompletedJobs(batch);
           completed.forEach(id => allCompleted.add(id));
-          ++numRequests;
-        }
-        if (numRequests === 0) {
-          const final = await this.broker.getCompletedJobs([]);
-          final.forEach(id => allCompleted.add(id));
-        }
+        });
       } catch (err) {
         this.log.error(`Error thrown when requesting completed job notifications from the broker`, err);
       }
@@ -292,6 +288,8 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
         }
       } else if (result.status === 'rejected') {
         return { success: false, reason: result.reason };
+      } else if (result.status === 'aborted') {
+        return { success: false, reason: 'Aborted' };
       } else {
         throw new Error(`Unexpected proving job status ${result.status}`);
       }
@@ -352,12 +350,8 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
       .map(id => this.jobs.get(id)!)
       .filter(x => x !== undefined);
     const totalJobsToRetrieve = toBeRetrieved.length;
-    let totalJobsRetrieved = 0;
-    while (toBeRetrieved.length > 0) {
-      const slice = toBeRetrieved.splice(0, MAX_CONCURRENT_JOB_SETTLED_REQUESTS);
-      const results = await Promise.all(slice.map(job => processJob(job!)));
-      totalJobsRetrieved += results.filter(x => x).length;
-    }
+    const results = await asyncPool(MAX_CONCURRENT_JOB_SETTLED_REQUESTS, toBeRetrieved, job => processJob(job));
+    const totalJobsRetrieved = results.filter(x => x).length;
     if (totalJobsToRetrieve > 0) {
       this.log.verbose(
         `Successfully retrieved ${totalJobsRetrieved} of ${totalJobsToRetrieve} jobs that should be ready, total ready jobs is now: ${this.jobsToRetrieve.size}`,
@@ -402,7 +396,7 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
     inputs: AvmCircuitInputs,
     signal?: AbortSignal,
     epochNumber?: EpochNumber,
-  ): Promise<RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>> {
+  ): Promise<RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS>> {
     return this.enqueueJob(
       this.generateId(ProvingRequestType.PUBLIC_VM, inputs, epochNumber),
       ProvingRequestType.PUBLIC_VM,

package/src/proving_broker/config.ts

@@ -4,10 +4,11 @@ import {
   booleanConfigHelper,
   getDefaultConfig,
   numberConfigHelper,
+  optionalNumberConfigHelper,
 } from '@aztec/foundation/config';
 import { pickConfigMappings } from '@aztec/foundation/config';
-import { type DataStoreConfig, dataConfigMappings } from '@aztec/kv-store/config';
 import { type ChainConfig, chainConfigMappings } from '@aztec/stdlib/config';
+import { type DataStoreConfig, dataConfigMappings } from '@aztec/stdlib/kv-store';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 
 import { z } from 'zod';
@@ -73,7 +74,7 @@ export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig>
   },
   proverBrokerStoreMapSizeKb: {
     env: 'PROVER_BROKER_STORE_MAP_SIZE_KB',
-    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    ...optionalNumberConfigHelper(),
     description: "The size of the prover broker's database. Will override the dataStoreMapSizeKb if set.",
   },
   proverBrokerDebugReplayEnabled: {

package/src/proving_broker/index.ts

@@ -7,3 +7,4 @@ export * from './proving_broker_database/persisted.js';
 export * from './proof_store/index.js';
 export * from './factory.js';
 export * from './config.js';
+export { BrokerCircuitProverFacade } from './broker_prover_facade.js';