feat: Process blocks in parallel during epoch proving

yarn-project/circuit-types/src/interfaces/epoch-prover.ts

@@ -1,32 +1,20 @@
-import { type Fr, type Proof, type RootRollupPublicInputs } from '@aztec/circuits.js';
+import { type Fr, type Header, type Proof, type RootRollupPublicInputs } from '@aztec/circuits.js';
 
 import { type L2Block } from '../l2_block.js';
 import { type BlockBuilder } from './block-builder.js';
 
-/**
- * Coordinates the proving of an entire epoch.
- *
- * Expected usage:
- * ```
- * startNewEpoch
- * foreach block {
- *   addNewBlock
- *   foreach tx {
- *     addTx
- *   }
- *   setBlockCompleted
- * }
- * finaliseEpoch
- * ```
- */
-export interface EpochProver extends BlockBuilder {
+/** Coordinates the proving of an entire epoch. */
+export interface EpochProver extends Omit<BlockBuilder, 'setBlockCompleted'> {
   /**
    * Starts a new epoch. Must be the first method to be called.
    * @param epochNumber - The epoch number.
    * @param totalNumBlocks - The total number of blocks expected in the epoch (must be at least one).
    **/
   startNewEpoch(epochNumber: number, totalNumBlocks: number): void;
 
+  /** Pads the block with empty txs if it hasn't reached the declared number of txs. */
+  setBlockCompleted(blockNumber: number, expectedBlockHeader?: Header): Promise<L2Block>;
+
   /** Pads the epoch with empty block roots if needed and blocks until proven. Throws if proving has failed. */
   finaliseEpoch(): Promise<{ publicInputs: RootRollupPublicInputs; proof: Proof }>;
 

yarn-project/circuit-types/src/interfaces/prover-client.ts

@@ -6,7 +6,6 @@ import { z } from 'zod';
 
 import { type TxHash } from '../tx/tx_hash.js';
 import { type EpochProver } from './epoch-prover.js';
-import { type MerkleTreeReadOperations } from './merkle_tree_operations.js';
 import { type ProvingJobConsumer } from './prover-broker.js';
 import { type ProvingJobStatus } from './proving-job.js';
 
@@ -105,7 +104,7 @@ export interface ProverCache {
  * Provides the ability to generate proofs and build rollups.
  */
 export interface EpochProverManager {
-  createEpochProver(db: MerkleTreeReadOperations, cache?: ProverCache): EpochProver;
+  createEpochProver(cache?: ProverCache): EpochProver;
 
   start(): Promise<void>;
 

yarn-project/circuit-types/src/interfaces/world_state.ts

@@ -25,10 +25,17 @@ export interface WorldStateSynchronizerStatus {
   syncedToL2Block: L2BlockId;
 }
 
-/**
- * Defines the interface for a world state synchronizer.
- */
-export interface WorldStateSynchronizer {
+/** Provides writeable forks of the world state at a given block number. */
+export interface ForkMerkleTreeOperations {
+  /** Forks the world state at the given block number, defaulting to the latest one. */
+  fork(block?: number): Promise<MerkleTreeWriteOperations>;
+
+  /** Gets a handle that allows reading the state as it was at the given block number. */
+  getSnapshot(blockNumber: number): MerkleTreeReadOperations;
+}
+
+/** Defines the interface for a world state synchronizer. */
+export interface WorldStateSynchronizer extends ForkMerkleTreeOperations {
   /**
    * Starts the synchronizer.
    * @returns A promise that resolves once the initial sync is completed.
@@ -53,19 +60,8 @@ export interface WorldStateSynchronizer {
    */
   syncImmediate(minBlockNumber?: number): Promise<number>;
 
-  /**
-   * Forks the current in-memory state based off the current committed state, and returns an instance that cannot modify the underlying data store.
-   */
-  fork(block?: number): Promise<MerkleTreeWriteOperations>;
-
   /**
    * Returns an instance of MerkleTreeAdminOperations that will not include uncommitted data.
    */
   getCommitted(): MerkleTreeReadOperations;
-
-  /**
-   * Returns a readonly instance of MerkleTreeAdminOperations where the state is as it was at the given block number
-   * @param block - The block number to look at
-   */
-  getSnapshot(block: number): MerkleTreeReadOperations;
 }

yarn-project/circuit-types/src/test/factories.ts

@@ -53,14 +53,7 @@ export function makeBloatedProcessedTx({
   privateOnly?: boolean;
 } = {}) {
   seed *= 0x1000; // Avoid clashing with the previous mock values if seed only increases by 1.
-
-  if (!header) {
-    if (db) {
-      header = db.getInitialHeader();
-    } else {
-      header = makeHeader(seed);
-    }
-  }
+  header ??= db?.getInitialHeader() ?? makeHeader(seed);
 
   const txConstantData = TxConstantData.empty();
   txConstantData.historicalHeader = header;

yarn-project/end-to-end/src/composed/integration_l1_publisher.test.ts

@@ -25,6 +25,7 @@ import { OutboxAbi, RollupAbi } from '@aztec/l1-artifacts';
 import { SHA256Trunc, StandardTree } from '@aztec/merkle-tree';
 import { getVKTreeRoot } from '@aztec/noir-protocol-circuits-types';
 import { protocolContractTreeRoot } from '@aztec/protocol-contracts';
+import { LightweightBlockBuilder } from '@aztec/prover-client/block-builder';
 import { L1Publisher } from '@aztec/sequencer-client';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 import {
@@ -52,7 +53,6 @@ import {
 } from 'viem';
 import { type PrivateKeyAccount, privateKeyToAccount } from 'viem/accounts';
 
-import { LightweightBlockBuilder } from '../../../sequencer-client/src/block_builder/light.js';
 import { sendL1ToL2Message } from '../fixtures/l1_to_l2_messaging.js';
 import { setupL1Contracts } from '../fixtures/utils.js';
 

yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts

@@ -269,6 +269,7 @@ export class FullProverTest {
       proverAgentCount: 2,
       publisherPrivateKey: `0x${proverNodePrivateKey!.toString('hex')}`,
       proverNodeMaxPendingJobs: 100,
+      proverNodeMaxParallelBlocksPerEpoch: 32,
       proverNodePollingIntervalMs: 100,
       quoteProviderBasisPointFee: 100,
       quoteProviderBondAmount: 1000n,

yarn-project/end-to-end/src/fixtures/utils.ts

@@ -682,6 +682,7 @@ export async function createAndSyncProverNode(
     proverAgentCount: 2,
     publisherPrivateKey: proverNodePrivateKey,
     proverNodeMaxPendingJobs: 10,
+    proverNodeMaxParallelBlocksPerEpoch: 32,
     proverNodePollingIntervalMs: 200,
     quoteProviderBasisPointFee: 100,
     quoteProviderBondAmount: 1000n,

yarn-project/foundation/package.json

@@ -11,6 +11,7 @@
     "./prettier": "./.prettierrc.json",
     "./abi": "./dest/abi/index.js",
     "./async-map": "./dest/async-map/index.js",
+    "./async-pool": "./dest/async-pool/index.js",
     "./aztec-address": "./dest/aztec-address/index.js",
     "./collection": "./dest/collection/index.js",
     "./config": "./dest/config/index.js",
@@ -163,4 +164,4 @@
   "engines": {
     "node": ">=18"
   }
-}
+}

yarn-project/foundation/src/async-pool/index.ts

@@ -0,0 +1,50 @@
+/*
+ * Adapted from https://github.com/rxaviers/async-pool/blob/1.x/lib/es6.js
+ *
+ * Copyright (c) 2017 Rafael Xavier de Souza http://rafael.xavier.blog.br
+ *
+ * Permission is hereby granted, free of charge, to any person
+ * obtaining a copy of this software and associated documentation
+ * files (the "Software"), to deal in the Software without
+ * restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following
+ * conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+ * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+ * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+ * OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+/** Executes the given async function over the iterable, up to a determined number of promises in parallel. */
+export function asyncPool<T, R>(poolLimit: number, iterable: T[], iteratorFn: (item: T, iterable: T[]) => Promise<R>) {
+  let i = 0;
+  const ret: Promise<R>[] = [];
+  const executing: Set<Promise<R>> = new Set();
+  const enqueue = (): Promise<any> => {
+    if (i === iterable.length) {
+      return Promise.resolve();
+    }
+    const item = iterable[i++];
+    const p = Promise.resolve().then(() => iteratorFn(item, iterable));
+    ret.push(p);
+    executing.add(p);
+    const clean = () => executing.delete(p);
+    p.then(clean).catch(clean);
+    let r: Promise<any> = Promise.resolve();
+    if (executing.size >= poolLimit) {
+      r = Promise.race(executing);
+    }
+    return r.then(() => enqueue());
+  };
+  return enqueue().then(() => Promise.all(ret));
+}

yarn-project/foundation/src/collection/array.test.ts

@@ -1,4 +1,4 @@
-import { compactArray, removeArrayPaddingEnd, times, unique } from './array.js';
+import { compactArray, maxBy, removeArrayPaddingEnd, times, unique } from './array.js';
 
 describe('times', () => {
   it('should return an array with the result from all executions', () => {
@@ -61,3 +61,21 @@ describe('unique', () => {
     expect(unique([1n, 2n, 1n])).toEqual([1n, 2n]);
   });
 });
+
+describe('maxBy', () => {
+  it('returns the max value', () => {
+    expect(maxBy([1, 2, 3], x => x)).toEqual(3);
+  });
+
+  it('returns the first max value', () => {
+    expect(maxBy([{ a: 1 }, { a: 3, b: 1 }, { a: 3, b: 2 }], ({ a }) => a)).toEqual({ a: 3, b: 1 });
+  });
+
+  it('returns undefined for an empty array', () => {
+    expect(maxBy([], x => x)).toBeUndefined();
+  });
+
+  it('applies the mapping function', () => {
+    expect(maxBy([1, 2, 3], x => -x)).toEqual(1);
+  });
+});

yarn-project/foundation/src/collection/array.ts

@@ -75,6 +75,20 @@ export function times<T>(n: number, fn: (i: number) => T): T[] {
   return [...Array(n).keys()].map(i => fn(i));
 }
 
+/**
+ * Executes the given async function n times and returns the results in an array. Awaits each execution before starting the next one.
+ * @param n - How many times to repeat.
+ * @param fn - Mapper from index to value.
+ * @returns The array with the result from all executions.
+ */
+export async function timesAsync<T>(n: number, fn: (i: number) => Promise<T>): Promise<T[]> {
+  const results: T[] = [];
+  for (let i = 0; i < n; i++) {
+    results.push(await fn(i));
+  }
+  return results;
+}
+
 /**
  * Returns the serialized size of all non-empty items in an array.
  * @param arr - Array
@@ -121,3 +135,13 @@ export function areArraysEqual<T>(a: T[], b: T[], eq: (a: T, b: T) => boolean =
   }
   return true;
 }
+
+/**
+ * Returns the element of the array that has the maximum value of the given function.
+ * In case of a tie, returns the first element with the maximum value.
+ * @param arr - The array.
+ * @param fn - The function to get the value to compare.
+ */
+export function maxBy<T>(arr: T[], fn: (x: T) => number): T | undefined {
+  return arr.reduce((max, x) => (fn(x) > fn(max) ? x : max), arr[0]);
+}

yarn-project/foundation/src/config/env_var.ts

@@ -115,6 +115,7 @@ export type EnvVar =
   | 'PROVER_JOB_TIMEOUT_MS'
   | 'PROVER_NODE_POLLING_INTERVAL_MS'
   | 'PROVER_NODE_MAX_PENDING_JOBS'
+  | 'PROVER_NODE_MAX_PARALLEL_BLOCKS_PER_EPOCH'
   | 'PROVER_PUBLISH_RETRY_INTERVAL_MS'
   | 'PROVER_PUBLISHER_PRIVATE_KEY'
   | 'PROVER_REAL_PROOFS'

yarn-project/prover-client/package.json

@@ -4,6 +4,7 @@
   "type": "module",
   "exports": {
     ".": "./dest/index.js",
+    "./block-builder": "./dest/block_builder/index.js",
     "./broker": "./dest/proving_broker/index.js",
     "./prover-agent": "./dest/prover-agent/index.js",
     "./orchestrator": "./dest/orchestrator/index.js",
@@ -103,4 +104,4 @@
   "engines": {
     "node": ">=18"
   }
-}
+}

yarn-project/sequencer-client/src/block_builder/index.ts → yarn-project/prover-client/src/block_builder/index.ts

@@ -1,6 +1,5 @@
 import { type BlockBuilder, type MerkleTreeReadOperations } from '@aztec/circuit-types';
 
-export * from './orchestrator.js';
 export * from './light.js';
 export interface BlockBuilderFactory {
   create(db: MerkleTreeReadOperations): BlockBuilder;

yarn-project/sequencer-client/src/block_builder/light.test.ts → yarn-project/prover-client/src/block_builder/light.test.ts

@@ -47,18 +47,18 @@ import {
   getVKTreeRoot,
 } from '@aztec/noir-protocol-circuits-types';
 import { protocolContractTreeRoot } from '@aztec/protocol-contracts';
+import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
+import { type MerkleTreeAdminDatabase, NativeWorldStateService } from '@aztec/world-state';
+
+import { jest } from '@jest/globals';
+
 import {
   buildBaseRollupHints,
   buildHeaderFromCircuitOutputs,
   getRootTreeSiblingPath,
   getSubtreeSiblingPath,
   getTreeSnapshot,
-} from '@aztec/prover-client/helpers';
-import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
-import { type MerkleTreeAdminDatabase, NativeWorldStateService } from '@aztec/world-state';
-
-import { jest } from '@jest/globals';
-
+} from '../orchestrator/block-building-helpers.js';
 import { LightweightBlockBuilder } from './light.js';
 
 jest.setTimeout(50_000);

yarn-project/sequencer-client/src/block_builder/light.ts → yarn-project/prover-client/src/block_builder/light.ts

@@ -1,4 +1,3 @@
-import { createDebugLogger } from '@aztec/aztec.js';
 import {
   type BlockBuilder,
   L2Block,
@@ -9,14 +8,20 @@ import {
 } from '@aztec/circuit-types';
 import { Fr, type GlobalVariables, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP } from '@aztec/circuits.js';
 import { padArrayEnd } from '@aztec/foundation/collection';
+import { createDebugLogger } from '@aztec/foundation/log';
 import { getVKTreeRoot } from '@aztec/noir-protocol-circuits-types';
 import { protocolContractTreeRoot } from '@aztec/protocol-contracts';
-import { buildBaseRollupHints, buildHeaderAndBodyFromTxs, getTreeSnapshot } from '@aztec/prover-client/helpers';
 import { type TelemetryClient } from '@aztec/telemetry-client';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 
 import { inspect } from 'util';
 
+import {
+  buildBaseRollupHints,
+  buildHeaderAndBodyFromTxs,
+  getTreeSnapshot,
+} from '../orchestrator/block-building-helpers.js';
+
 /**
  * Builds a block and its header from a set of processed tx without running any circuits.
  */
@@ -90,3 +95,23 @@ export class LightweightBlockBuilderFactory {
     return new LightweightBlockBuilder(db, this.telemetry ?? new NoopTelemetryClient());
   }
 }
+
+/**
+ * Creates a block builder under the hood with the given txs and messages and creates a block.
+ * Automatically adds padding txs to get to a minimum of 2 txs in the block.
+ * @param db - A db fork to use for block building.
+ */
+export async function buildBlock(
+  txs: ProcessedTx[],
+  globalVariables: GlobalVariables,
+  l1ToL2Messages: Fr[],
+  db: MerkleTreeWriteOperations,
+  telemetry: TelemetryClient = new NoopTelemetryClient(),
+) {
+  const builder = new LightweightBlockBuilder(db, telemetry);
+  await builder.startNewBlock(Math.max(txs.length, 2), globalVariables, l1ToL2Messages);
+  for (const tx of txs) {
+    await builder.addNewTx(tx);
+  }
+  return await builder.setBlockCompleted();
+}

yarn-project/prover-client/src/index.ts

@@ -1,6 +1,5 @@
 export { EpochProverManager } from '@aztec/circuit-types';
 
-export * from './tx-prover/tx-prover.js';
+export * from './prover-client/index.js';
 export * from './config.js';
-export * from './tx-prover/factory.js';
 export * from './proving_broker/prover_cache/memory.js';