@aztec-labs/prover-client 6.0.0-nightly.20260829

package/src/mocks/test_context.ts

@@ -0,0 +1,477 @@
+import type { BBProverConfig } from '@aztec-labs/bb-prover';
+import { TestCircuitProver } from '@aztec-labs/bb-prover';
+import { BlockNumber, CheckpointNumber } from '@aztec-labs/foundation/branded-types';
+import { times, timesAsync } from '@aztec-labs/foundation/collection';
+import { Fr } from '@aztec-labs/foundation/curves/bn254';
+import type { Logger } from '@aztec-labs/foundation/log';
+import { SerialQueue } from '@aztec-labs/foundation/queue';
+import type { FieldsOf } from '@aztec-labs/foundation/types';
+import { getVKTreeRoot } from '@aztec-labs/noir-protocol-circuits-types/vk-tree';
+import { ProtocolContractsList } from '@aztec-labs/protocol-contracts';
+import { computeFeePayerBalanceLeafSlot } from '@aztec-labs/protocol-contracts/fee-juice';
+import { PublicDataWrite } from '@aztec-labs/stdlib/avm';
+import { AztecAddress } from '@aztec-labs/stdlib/aztec-address';
+import type { Checkpoint } from '@aztec-labs/stdlib/checkpoint';
+import type { MerkleTreeWriteOperations, ServerCircuitProver } from '@aztec-labs/stdlib/interfaces/server';
+import type { CheckpointConstantData } from '@aztec-labs/stdlib/rollup';
+import { mockProcessedTx } from '@aztec-labs/stdlib/testing';
+import { MerkleTreeId, PublicDataTreeLeaf } from '@aztec-labs/stdlib/trees';
+import {
+  type BlockHeader,
+  type GlobalVariables,
+  type ProcessedTx,
+  StateReference,
+  TreeSnapshots,
+} from '@aztec-labs/stdlib/tx';
+import type { GenesisData } from '@aztec-labs/stdlib/world-state';
+import type { MerkleTreeAdminDatabase } from '@aztec-labs/world-state';
+import { NativeWorldStateService } from '@aztec-labs/world-state/native';
+import { promises as fs } from 'fs';
+
+import { LightweightCheckpointBuilder } from '../light/lightweight_checkpoint_builder.js';
+import {
+  buildFinalBlobChallenges,
+  getTreeSnapshot,
+  insertSideEffects,
+} from '../orchestrator/block-building-helpers.js';
+import { BrokerCircuitProverFacade } from '../proving_broker/broker_prover_facade.js';
+import { TestBroker } from '../test/mock_prover.js';
+import { getEnvironmentConfig, getSimulator, makeCheckpointConstants, makeGlobals } from './fixtures.js';
+
+/**
+ * Builds a started `SerialQueue` for use as an orchestrator's deferred-job queue in
+ * tests. Production wires a single shared queue via `ProverClient`; tests that construct
+ * orchestrators directly use one of these per orchestrator (or share one).
+ */
+export function makeTestDeferredJobQueue(concurrency = 10): SerialQueue {
+  const queue = new SerialQueue();
+  queue.start(concurrency);
+  return queue;
+}
+
+export class TestContext {
+  private headers: Map<number, BlockHeader> = new Map();
+  private checkpoints: Checkpoint[] = [];
+  private checkpointOutHashes: Fr[] = [];
+  private nextCheckpointIndex = 0;
+  private nextCheckpointNumber = CheckpointNumber(1);
+  private nextBlockNumber = 1;
+  private epochNumber = 1;
+  private feePayerBalance: Fr;
+  // Running inbox rolling hash across the checkpoints built by this context (never resets: the protocol threads it
+  // from genesis). Each checkpoint's builder starts from it, and it advances by the checkpoint's messages, so a
+  // checkpoint built after a message-carrying one commits the correct continuation in its header.
+  private currentInboxRollingHash = Fr.ZERO;
+
+  constructor(
+    public worldState: MerkleTreeAdminDatabase,
+    public prover: ServerCircuitProver,
+    public broker: TestBroker,
+    public brokerProverFacade: BrokerCircuitProverFacade,
+    private feePayer: AztecAddress,
+    initialFeePayerBalance: Fr,
+    private directoriesToCleanup: string[],
+    private logger: Logger,
+  ) {
+    this.feePayerBalance = initialFeePayerBalance;
+  }
+
+  static async new(
+    logger: Logger,
+    {
+      proverCount = 4,
+      createProver = async (bbConfig: BBProverConfig) => new TestCircuitProver(await getSimulator(bbConfig, logger)),
+    }: {
+      proverCount?: number;
+      createProver?: (bbConfig: BBProverConfig) => Promise<ServerCircuitProver>;
+    } = {},
+  ) {
+    const directoriesToCleanup: string[] = [];
+
+    const feePayer = AztecAddress.fromNumberUnsafe(42222);
+    const initialFeePayerBalance = new Fr(10n ** 20n);
+    const feePayerSlot = await computeFeePayerBalanceLeafSlot(feePayer);
+    const genesis: GenesisData = {
+      prefilledPublicData: [new PublicDataTreeLeaf(feePayerSlot, initialFeePayerBalance)],
+      genesisTimestamp: 0n,
+    };
+
+    // Separated dbs for public processor and prover - see public_processor for context
+    const ws = await NativeWorldStateService.tmp(/*cleanupTmpDir=*/ true, genesis);
+
+    let localProver: ServerCircuitProver;
+    const config = await getEnvironmentConfig(logger);
+    if (!config) {
+      localProver = new TestCircuitProver();
+    } else {
+      const bbConfig: BBProverConfig = {
+        acvmBinaryPath: config.expectedAcvmPath,
+        acvmWorkingDirectory: config.acvmWorkingDirectory,
+        bbBinaryPath: config.expectedBBPath,
+        bbWorkingDirectory: config.bbWorkingDirectory,
+        bbSkipCleanup: config.bbSkipCleanup,
+        numConcurrentIVCVerifiers: 8,
+        bbIVCConcurrency: 1,
+        bbChonkVerifyMaxBatch: 16,
+        bbChonkVerifyConcurrency: 6,
+      };
+      localProver = await createProver(bbConfig);
+    }
+
+    if (config?.directoryToCleanup && !config.bbSkipCleanup) {
+      directoriesToCleanup.push(config.directoryToCleanup);
+    }
+
+    const broker = new TestBroker(proverCount, localProver);
+    const facade = new BrokerCircuitProverFacade(broker);
+
+    await broker.start();
+    facade.start();
+
+    return new this(ws, localProver, broker, facade, feePayer, initialFeePayerBalance, directoriesToCleanup, logger);
+  }
+
+  public getFork() {
+    return this.worldState.fork();
+  }
+
+  async cleanup() {
+    await this.brokerProverFacade.stop();
+    await this.broker.stop();
+    await this.worldState.close();
+    for (const dir of this.directoriesToCleanup.filter(x => x !== '')) {
+      try {
+        await fs.rm(dir, { recursive: true, force: true, maxRetries: 3 });
+      } catch (err) {
+        this.logger.warn(`Failed to delete tmp directory $dir}: ${err}`);
+      }
+    }
+  }
+
+  public startNewEpoch() {
+    this.checkpoints = [];
+    this.checkpointOutHashes = [];
+    this.nextCheckpointIndex = 0;
+    this.epochNumber++;
+  }
+
+  // Return blob fields of all checkpoints in the epoch.
+  public getBlobFields() {
+    return this.checkpoints.map(checkpoint => checkpoint.toBlobFields());
+  }
+
+  public async getFinalBlobChallenges() {
+    const blobFields = this.getBlobFields();
+    return await buildFinalBlobChallenges(blobFields);
+  }
+
+  public async makeCheckpoint(
+    numBlocks: number,
+    {
+      numTxsPerBlock = 0,
+      numL1ToL2Messages = 0,
+      makeProcessedTxOpts = () => ({}),
+      ...constantOpts
+    }: {
+      numTxsPerBlock?: number | number[];
+      numL1ToL2Messages?: number;
+      makeProcessedTxOpts?: (
+        blockGlobalVariables: GlobalVariables,
+        txIndex: number,
+      ) => Partial<Parameters<typeof mockProcessedTx>[0]>;
+    } & Partial<FieldsOf<CheckpointConstantData>> = {},
+  ) {
+    if (numBlocks === 0) {
+      throw new Error(
+        'Cannot make a checkpoint with 0 blocks. Crate an empty block (numTxsPerBlock = 0) if there are no txs.',
+      );
+    }
+
+    const checkpointIndex = this.nextCheckpointIndex++;
+    const checkpointNumber = this.nextCheckpointNumber;
+    this.nextCheckpointNumber++;
+    const slotNumber = checkpointNumber * 15; // times an arbitrary number to make it different to the checkpoint number
+
+    const constants = makeCheckpointConstants(slotNumber, constantOpts);
+
+    const fork = await this.worldState.fork();
+
+    // Build l1 to l2 messages. Appended unpadded at compact indices; the mock assigns them all to
+    // the checkpoint's first block, matching how the per-block driver slices them.
+    const l1ToL2Messages = times(numL1ToL2Messages, i => new Fr(slotNumber * 100 + i));
+    await fork.appendLeaves(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l1ToL2Messages);
+    const newL1ToL2Snapshot = await getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, fork);
+
+    const startBlockNumber = this.nextBlockNumber;
+    const previousBlockHeader = this.getBlockHeader(BlockNumber(startBlockNumber - 1));
+    // All blocks in the same slot/checkpoint share the same timestamp.
+    const timestamp = BigInt(slotNumber * 26);
+
+    // Build global variables.
+    const blockGlobalVariables = times(numBlocks, i =>
+      makeGlobals(startBlockNumber + i, slotNumber, {
+        coinbase: constants.coinbase,
+        feeRecipient: constants.feeRecipient,
+        gasFees: constants.gasFees,
+        timestamp,
+      }),
+    );
+    this.nextBlockNumber += numBlocks;
+
+    // Build txs.
+    let totalTxs = 0;
+    const blockEndStates: StateReference[] = [];
+    const blockTxs = await timesAsync(numBlocks, async blockIndex => {
+      const txIndexOffset = totalTxs;
+      const numTxs = typeof numTxsPerBlock === 'number' ? numTxsPerBlock : numTxsPerBlock[blockIndex];
+      totalTxs += numTxs;
+      const txs = await timesAsync(numTxs, txIndex =>
+        this.makeProcessedTx({
+          seed: (txIndexOffset + txIndex + 1) * 321 + (checkpointIndex + 1) * 123456 + this.epochNumber * 0x99999,
+          globalVariables: blockGlobalVariables[blockIndex],
+          anchorBlockHeader: previousBlockHeader,
+          newL1ToL2Snapshot,
+          ...makeProcessedTxOpts(blockGlobalVariables[blockIndex], txIndexOffset + txIndex),
+        }),
+      );
+
+      // Insert side effects into the trees.
+      const endState = await this.updateTrees(txs, fork);
+      blockEndStates.push(endState);
+
+      return txs;
+    });
+
+    const cleanFork = await this.worldState.fork();
+    const previousCheckpointOutHashes = this.checkpointOutHashes;
+    const startInboxRollingHash = this.currentInboxRollingHash;
+    const builder = LightweightCheckpointBuilder.startNewCheckpoint(
+      checkpointNumber,
+      { ...constants, timestamp },
+      previousCheckpointOutHashes,
+      startInboxRollingHash,
+      cleanFork,
+    );
+
+    // Add tx effects to db and build block headers.
+    const blocks = [];
+    for (let i = 0; i < numBlocks; i++) {
+      const txs = blockTxs[i];
+      const state = blockEndStates[i];
+
+      const { block } = await builder.addBlock(blockGlobalVariables[i], txs, i === 0 ? l1ToL2Messages : [], {
+        expectedEndState: state,
+        insertTxsEffects: true,
+      });
+
+      const header = block.header;
+      this.headers.set(block.number, header);
+
+      const blockMsgs = block.indexWithinCheckpoint === 0 ? l1ToL2Messages : [];
+      await this.worldState.handleL2BlockAndMessages(block, blockMsgs);
+
+      blocks.push({ header, txs });
+    }
+
+    const checkpoint = await builder.completeCheckpoint();
+    this.checkpoints.push(checkpoint);
+    this.checkpointOutHashes.push(checkpoint.getCheckpointOutHash());
+    this.currentInboxRollingHash = checkpoint.header.inboxRollingHash;
+
+    return {
+      constants,
+      checkpoint,
+      header: checkpoint.header,
+      blocks,
+      l1ToL2Messages,
+      previousBlockHeader,
+      startInboxRollingHash,
+    };
+  }
+
+  /**
+   * Like {@link makeCheckpoint} but distributes the L1-to-L2 message bundle across the checkpoint's blocks
+   * (streaming Inbox): `l1ToL2MessagesPerBlock[i]` is block `i`'s own message slice,
+   * appended at compact indices in insertion order. Lets tests exercise checkpoints whose messages span more
+   * than one block, including a non-first block carrying a bundle — the single-block-per-checkpoint
+   * `makeCheckpoint` puts every message in the first block. `numTxsPerBlock` defaults to 1 because the builder
+   * rejects a non-first block with neither txs nor messages; a zero-tx entry whose slice is non-empty is valid
+   * and produces a message-only block.
+   */
+  public async makeCheckpointWithMessagesPerBlock(
+    l1ToL2MessagesPerBlock: Fr[][],
+    {
+      numTxsPerBlock = 1,
+      makeProcessedTxOpts = () => ({}),
+      ...constantOpts
+    }: {
+      numTxsPerBlock?: number | number[];
+      makeProcessedTxOpts?: (
+        blockGlobalVariables: GlobalVariables,
+        txIndex: number,
+      ) => Partial<Parameters<typeof mockProcessedTx>[0]>;
+    } & Partial<FieldsOf<CheckpointConstantData>> = {},
+  ) {
+    const numBlocks = l1ToL2MessagesPerBlock.length;
+    if (numBlocks === 0) {
+      throw new Error('Cannot make a checkpoint with 0 blocks.');
+    }
+
+    const checkpointIndex = this.nextCheckpointIndex++;
+    const checkpointNumber = this.nextCheckpointNumber;
+    this.nextCheckpointNumber++;
+    const slotNumber = checkpointNumber * 15;
+
+    const constants = makeCheckpointConstants(slotNumber, constantOpts);
+    const l1ToL2Messages = l1ToL2MessagesPerBlock.flat();
+
+    const fork = await this.worldState.fork();
+
+    const startBlockNumber = this.nextBlockNumber;
+    const previousBlockHeader = this.getBlockHeader(BlockNumber(startBlockNumber - 1));
+    const timestamp = BigInt(slotNumber * 26);
+
+    const blockGlobalVariables = times(numBlocks, i =>
+      makeGlobals(startBlockNumber + i, slotNumber, {
+        coinbase: constants.coinbase,
+        feeRecipient: constants.feeRecipient,
+        gasFees: constants.gasFees,
+        timestamp,
+      }),
+    );
+    this.nextBlockNumber += numBlocks;
+
+    // Build txs per block. Append the block's message slice to the fork before its txs so the per-block end
+    // state matches the builder's (message-tree and tx-effect trees are independent, so append order does not
+    // matter, but the cumulative message tree must reflect the slices already inserted).
+    let totalTxs = 0;
+    const blockEndStates: StateReference[] = [];
+    const blockTxs = await timesAsync(numBlocks, async blockIndex => {
+      await fork.appendLeaves(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l1ToL2MessagesPerBlock[blockIndex]);
+      const newL1ToL2Snapshot = await getTreeSnapshot(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, fork);
+
+      const txIndexOffset = totalTxs;
+      const numTxs = typeof numTxsPerBlock === 'number' ? numTxsPerBlock : numTxsPerBlock[blockIndex];
+      totalTxs += numTxs;
+      const txs = await timesAsync(numTxs, txIndex =>
+        this.makeProcessedTx({
+          seed: (txIndexOffset + txIndex + 1) * 321 + (checkpointIndex + 1) * 123456 + this.epochNumber * 0x99999,
+          globalVariables: blockGlobalVariables[blockIndex],
+          anchorBlockHeader: previousBlockHeader,
+          newL1ToL2Snapshot,
+          ...makeProcessedTxOpts(blockGlobalVariables[blockIndex], txIndexOffset + txIndex),
+        }),
+      );
+
+      const endState = await this.updateTrees(txs, fork);
+      blockEndStates.push(endState);
+
+      return txs;
+    });
+
+    const cleanFork = await this.worldState.fork();
+    const previousCheckpointOutHashes = this.checkpointOutHashes;
+    const startInboxRollingHash = this.currentInboxRollingHash;
+    const builder = LightweightCheckpointBuilder.startNewCheckpoint(
+      checkpointNumber,
+      { ...constants, timestamp },
+      previousCheckpointOutHashes,
+      startInboxRollingHash,
+      cleanFork,
+    );
+
+    const blocks = [];
+    for (let i = 0; i < numBlocks; i++) {
+      const txs = blockTxs[i];
+      const state = blockEndStates[i];
+
+      const { block } = await builder.addBlock(blockGlobalVariables[i], txs, l1ToL2MessagesPerBlock[i], {
+        expectedEndState: state,
+        insertTxsEffects: true,
+      });
+
+      const header = block.header;
+      this.headers.set(block.number, header);
+
+      await this.worldState.handleL2BlockAndMessages(block, l1ToL2MessagesPerBlock[i]);
+
+      blocks.push({ header, txs });
+    }
+
+    const checkpoint = await builder.completeCheckpoint();
+    this.checkpoints.push(checkpoint);
+    this.checkpointOutHashes.push(checkpoint.getCheckpointOutHash());
+    this.currentInboxRollingHash = checkpoint.header.inboxRollingHash;
+
+    return {
+      constants,
+      checkpoint,
+      header: checkpoint.header,
+      blocks,
+      l1ToL2Messages,
+      l1ToL2MessagesPerBlock,
+      previousBlockHeader,
+      startInboxRollingHash,
+    };
+  }
+
+  private async makeProcessedTx(opts: Parameters<typeof mockProcessedTx>[0] = {}): Promise<ProcessedTx> {
+    const tx = await mockProcessedTx({
+      vkTreeRoot: getVKTreeRoot(),
+      protocolContracts: ProtocolContractsList,
+      feePayer: this.feePayer,
+      ...opts,
+    });
+
+    this.feePayerBalance = new Fr(this.feePayerBalance.toBigInt() - tx.txEffect.transactionFee.toBigInt());
+
+    const feePayerSlot = await computeFeePayerBalanceLeafSlot(this.feePayer);
+    const feePaymentPublicDataWrite = new PublicDataWrite(feePayerSlot, this.feePayerBalance);
+    tx.txEffect.publicDataWrites[0] = feePaymentPublicDataWrite;
+    if (tx.avmProvingRequest) {
+      tx.avmProvingRequest.inputs.publicInputs.accumulatedData.publicDataWrites[0] = feePaymentPublicDataWrite;
+    }
+
+    return tx;
+  }
+
+  private getBlockHeader(blockNumber: BlockNumber): BlockHeader {
+    if (Number(blockNumber) > 0 && Number(blockNumber) >= this.nextBlockNumber) {
+      throw new Error(`Block header not built for block number ${blockNumber}.`);
+    }
+    return Number(blockNumber) === 0
+      ? this.worldState.getCommitted().getInitialHeader()
+      : this.headers.get(Number(blockNumber))!;
+  }
+
+  private async updateTrees(txs: ProcessedTx[], fork: MerkleTreeWriteOperations) {
+    let startStateReference = await fork.getStateReference();
+    let endStateReference = startStateReference;
+    for (const tx of txs) {
+      await insertSideEffects(tx, fork);
+      endStateReference = await fork.getStateReference();
+
+      if (tx.avmProvingRequest) {
+        // Update the trees in the avm public inputs so that the proof won't fail.
+        const l1ToL2MessageTree = tx.avmProvingRequest.inputs.publicInputs.startTreeSnapshots.l1ToL2MessageTree;
+        tx.avmProvingRequest.inputs.publicInputs.startTreeSnapshots = new TreeSnapshots(
+          l1ToL2MessageTree,
+          startStateReference.partial.noteHashTree,
+          startStateReference.partial.nullifierTree,
+          startStateReference.partial.publicDataTree,
+        );
+
+        tx.avmProvingRequest.inputs.publicInputs.endTreeSnapshots = new TreeSnapshots(
+          l1ToL2MessageTree,
+          endStateReference.partial.noteHashTree,
+          endStateReference.partial.nullifierTree,
+          endStateReference.partial.publicDataTree,
+        );
+      }
+
+      startStateReference = endStateReference;
+    }
+
+    return endStateReference;
+  }
+}