fix(bb-prover): create structure for AVM vk (#8233)

Apologies for duplicating code! I tried putting a generic on the "base"
classes, but (1) generics don't play well with static methods (e.g.,
fromBuffer) and (2) you still need to pass the value for the VK size (on
top of the type). I think most of this duplication can be avoided if you
just accept some type unsafety and save things as `Fr[]` instead of
tuples with size.

PS: There might be still work to do to align the "num public inputs" etc
indices, and the vk hash.

barretenberg/cpp/src/barretenberg/vm/aztec_constants.hpp

@@ -33,6 +33,7 @@
 #define HEADER_LENGTH 24
 #define PUBLIC_CIRCUIT_PUBLIC_INPUTS_LENGTH 691
 #define PUBLIC_CONTEXT_INPUTS_LENGTH 42
+#define AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS 66
 #define SENDER_SELECTOR 0
 #define ADDRESS_SELECTOR 1
 #define STORAGE_ADDRESS_SELECTOR 1

noir-projects/noir-protocol-circuits/crates/types/src/constants.nr

@@ -290,6 +290,12 @@ global NESTED_RECURSIVE_PROOF_LENGTH = 439;
 global TUBE_PROOF_LENGTH = RECURSIVE_PROOF_LENGTH; // in the future these can differ
 
 global VERIFICATION_KEY_LENGTH_IN_FIELDS = 128;
+// VK is composed of
+// - circuit size encoded as a fr field element (32 bytes)
+// - num of inputs encoded as a fr field element (32 bytes)
+// - 16 affine elements (curve base field fq) encoded as fr elements takes (16 * 4 * 32 bytes)
+// 16 above refers to the constant AvmFlavor::NUM_PRECOMPUTED_ENTITIES
+global AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS = 2 + 16 * 4;
 
 /**
  * Enumerate the hash_indices which are used for pedersen hashing.

yarn-project/bb-prover/src/avm_proving.test.ts

@@ -1,5 +1,6 @@
 import {
   AvmCircuitInputs,
+  AvmVerificationKeyData,
   AztecAddress,
   ContractStorageRead,
   ContractStorageUpdateRequest,
@@ -49,7 +50,7 @@ import path from 'path';
 import { PublicSideEffectTrace } from '../../simulator/src/public/side_effect_trace.js';
 import { SerializableContractInstance } from '../../types/src/contracts/contract_instance.js';
 import { type BBSuccess, BB_RESULT, generateAvmProof, verifyAvmProof } from './bb/execute.js';
-import { extractVkData } from './verification_key/verification_key_data.js';
+import { extractAvmVkData } from './verification_key/verification_key_data.js';
 
 const TIMEOUT = 60_000;
 const TIMESTAMP = new Fr(99833);
@@ -279,18 +280,10 @@ const proveAndVerifyAvmTestContract = async (
   const proofRes = await generateAvmProof(bbPath, bbWorkingDirectory, avmCircuitInputs, logger);
   expect(proofRes.status).toEqual(BB_RESULT.SUCCESS);
 
-  // Then we test VK extraction.
+  // Then we test VK extraction and serialization.
   const succeededRes = proofRes as BBSuccess;
-  const verificationKey = await extractVkData(succeededRes.vkPath!);
-
-  // VK is composed of
-  // - circuit size encoded as a fr field element (32 bytes)
-  // - num of inputs encoded as a fr field element (32 bytes)
-  // - 16 affine elements (curve base field fq) encoded as fr elements takes (16 * 4 * 32 bytes)
-  // 16 above refers to the constant AvmFlavor::NUM_PRECOMPUTED_ENTITIES
-  // Total number of bytes = 2112
-  const NUM_PRECOMPUTED_ENTITIES = 16;
-  expect(verificationKey.keyAsBytes).toHaveLength(NUM_PRECOMPUTED_ENTITIES * 4 * 32 + 2 * 32);
+  const vkData = await extractAvmVkData(succeededRes.vkPath!);
+  AvmVerificationKeyData.fromBuffer(vkData.toBuffer());
 
   // Then we verify.
   const rawVkPath = path.join(succeededRes.vkPath!, 'vk');

yarn-project/bb-prover/src/prover/bb_prover.ts

@@ -1,6 +1,6 @@
 /* eslint-disable require-await */
 import {
-  type ProofAndVerificationKey,
+  type AvmProofAndVerificationKey,
   type PublicInputsAndRecursiveProof,
   type PublicKernelNonTailRequest,
   type PublicKernelTailRequest,
@@ -11,6 +11,7 @@ import { type CircuitProvingStats, type CircuitWitnessGenerationStats } from '@a
 import {
   AGGREGATION_OBJECT_LENGTH,
   type AvmCircuitInputs,
+  type AvmVerificationKeyData,
   type BaseOrMergeRollupPublicInputs,
   type BaseParityInputs,
   type BaseRollupInputs,
@@ -94,7 +95,7 @@ import type { ACVMConfig, BBConfig } from '../config.js';
 import { ProverInstrumentation } from '../instrumentation.js';
 import { PublicKernelArtifactMapping } from '../mappings/mappings.js';
 import { mapProtocolArtifactNameToCircuitName } from '../stats.js';
-import { extractVkData } from '../verification_key/verification_key_data.js';
+import { extractAvmVkData, extractVkData } from '../verification_key/verification_key_data.js';
 
 const logger = createDebugLogger('aztec:bb-prover');
 
@@ -202,7 +203,7 @@ export class BBNativeRollupProver implements ServerCircuitProver {
   @trackSpan('BBNativeRollupProver.getAvmProof', inputs => ({
     [Attributes.APP_CIRCUIT_NAME]: inputs.functionName,
   }))
-  public async getAvmProof(inputs: AvmCircuitInputs): Promise<ProofAndVerificationKey> {
+  public async getAvmProof(inputs: AvmCircuitInputs): Promise<AvmProofAndVerificationKey> {
     const proofAndVk = await this.createAvmProof(inputs);
     await this.verifyAvmProof(proofAndVk.proof, proofAndVk.verificationKey);
     return proofAndVk;
@@ -626,14 +627,14 @@ export class BBNativeRollupProver implements ServerCircuitProver {
     return provingResult;
   }
 
-  private async createAvmProof(input: AvmCircuitInputs): Promise<ProofAndVerificationKey> {
-    const operation = async (bbWorkingDirectory: string): Promise<ProofAndVerificationKey> => {
+  private async createAvmProof(input: AvmCircuitInputs): Promise<AvmProofAndVerificationKey> {
+    const operation = async (bbWorkingDirectory: string): Promise<AvmProofAndVerificationKey> => {
       const provingResult = await this.generateAvmProofWithBB(input, bbWorkingDirectory);
 
       const rawProof = await fs.readFile(provingResult.proofPath!);
       // TODO(https://github.com/AztecProtocol/aztec-packages/issues/6773): this VK data format is wrong.
       // In particular, the number of public inputs, etc will be wrong.
-      const verificationKey = await extractVkData(provingResult.vkPath!);
+      const verificationKey = await extractAvmVkData(provingResult.vkPath!);
       const proof = new Proof(rawProof, verificationKey.numPublicInputs);
 
       const circuitType = 'avm-circuit' as const;
@@ -765,7 +766,7 @@ export class BBNativeRollupProver implements ServerCircuitProver {
     return await this.verifyWithKey(verificationKey, proof);
   }
 
-  public async verifyAvmProof(proof: Proof, verificationKey: VerificationKeyData) {
+  public async verifyAvmProof(proof: Proof, verificationKey: AvmVerificationKeyData) {
     return await this.verifyWithKeyInternal(proof, verificationKey, verifyAvmProof);
   }
 
@@ -775,7 +776,7 @@ export class BBNativeRollupProver implements ServerCircuitProver {
 
   private async verifyWithKeyInternal(
     proof: Proof,
-    verificationKey: VerificationKeyData,
+    verificationKey: { keyAsBytes: Buffer },
     verificationFunction: VerificationFunction,
   ) {
     const operation = async (bbWorkingDirectory: string) => {

yarn-project/bb-prover/src/test/test_circuit_prover.ts

@@ -1,5 +1,5 @@
 import {
-  type ProofAndVerificationKey,
+  type AvmProofAndVerificationKey,
   type PublicInputsAndRecursiveProof,
   type PublicKernelNonTailRequest,
   type PublicKernelTailRequest,
@@ -8,6 +8,7 @@ import {
 } from '@aztec/circuit-types';
 import {
   type AvmCircuitInputs,
+  AvmVerificationKeyData,
   type BaseOrMergeRollupPublicInputs,
   type BaseParityInputs,
   type BaseRollupInputs,
@@ -475,12 +476,12 @@ export class TestCircuitProver implements ServerCircuitProver {
     );
   }
 
-  public async getAvmProof(_inputs: AvmCircuitInputs): Promise<ProofAndVerificationKey> {
+  public async getAvmProof(_inputs: AvmCircuitInputs): Promise<AvmProofAndVerificationKey> {
     // We can't simulate the AVM because we don't have enough context to do so (e.g., DBs).
     // We just return an empty proof and VK data.
     this.logger.debug('Skipping AVM simulation in TestCircuitProver.');
     await this.delay();
-    return { proof: makeEmptyProof(), verificationKey: VerificationKeyData.makeFake() };
+    return { proof: makeEmptyProof(), verificationKey: AvmVerificationKeyData.makeFake() };
   }
 
   private async delay(): Promise<void> {

yarn-project/bb-prover/src/verification_key/verification_key_data.ts

@@ -1,11 +1,15 @@
 import {
+  AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS,
+  AvmVerificationKeyAsFields,
+  AvmVerificationKeyData,
   Fr,
-  type VERIFICATION_KEY_LENGTH_IN_FIELDS,
+  VERIFICATION_KEY_LENGTH_IN_FIELDS,
   VerificationKeyAsFields,
   VerificationKeyData,
 } from '@aztec/circuits.js';
 import { type Tuple } from '@aztec/foundation/serialize';
 
+import { strict as assert } from 'assert';
 import * as fs from 'fs/promises';
 import * as path from 'path';
 
@@ -25,7 +29,25 @@ export async function extractVkData(vkDirectoryPath: string): Promise<Verificati
   const fields = fieldsJson.map(Fr.fromString);
   // The first item is the hash, this is not part of the actual VK
   const vkHash = fields[0];
+  assert(fields.length === VERIFICATION_KEY_LENGTH_IN_FIELDS, 'Invalid verification key length');
   const vkAsFields = new VerificationKeyAsFields(fields as Tuple<Fr, typeof VERIFICATION_KEY_LENGTH_IN_FIELDS>, vkHash);
   const vk = new VerificationKeyData(vkAsFields, rawBinary);
   return vk;
 }
+
+// TODO: This was adapted from the above function. A refactor might be needed.
+export async function extractAvmVkData(vkDirectoryPath: string): Promise<AvmVerificationKeyData> {
+  const [rawFields, rawBinary] = await Promise.all([
+    fs.readFile(path.join(vkDirectoryPath, VK_FIELDS_FILENAME), { encoding: 'utf-8' }),
+    fs.readFile(path.join(vkDirectoryPath, VK_FILENAME)),
+  ]);
+  const fieldsJson = JSON.parse(rawFields);
+  const fields = fieldsJson.map(Fr.fromString);
+  // The first item is the hash, this is not part of the actual VK
+  // TODO: is the above actually the case?
+  const vkHash = fields[0];
+  assert(fields.length === AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS, 'Invalid AVM verification key length');
+  const vkAsFields = new AvmVerificationKeyAsFields(fields, vkHash);
+  const vk = new AvmVerificationKeyData(vkAsFields, rawBinary);
+  return vk;
+}

yarn-project/circuit-types/src/interfaces/proving-job.ts

@@ -1,5 +1,6 @@
 import {
   type AvmCircuitInputs,
+  type AvmVerificationKeyData,
   type BaseOrMergeRollupPublicInputs,
   type BaseParityInputs,
   type BaseRollupInputs,
@@ -25,9 +26,9 @@ import {
 
 import type { PublicKernelNonTailRequest, PublicKernelTailRequest } from '../tx/processed_tx.js';
 
-export type ProofAndVerificationKey = {
+export type AvmProofAndVerificationKey = {
   proof: Proof;
-  verificationKey: VerificationKeyData;
+  verificationKey: AvmVerificationKeyData;
 };
 
 export type PublicInputsAndRecursiveProof<T> = {
@@ -133,7 +134,7 @@ export type ProvingRequest =
 
 export type ProvingRequestPublicInputs = {
   [ProvingRequestType.PRIVATE_KERNEL_EMPTY]: PublicInputsAndRecursiveProof<KernelCircuitPublicInputs>;
-  [ProvingRequestType.PUBLIC_VM]: ProofAndVerificationKey;
+  [ProvingRequestType.PUBLIC_VM]: AvmProofAndVerificationKey;
 
   [ProvingRequestType.PUBLIC_KERNEL_NON_TAIL]: PublicInputsAndRecursiveProof<PublicKernelCircuitPublicInputs>;
   [ProvingRequestType.PUBLIC_KERNEL_TAIL]: PublicInputsAndRecursiveProof<KernelCircuitPublicInputs>;

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

@@ -1,5 +1,5 @@
 import {
-  type ProofAndVerificationKey,
+  type AvmProofAndVerificationKey,
   type PublicInputsAndRecursiveProof,
   type PublicInputsAndTubeProof,
   type PublicKernelNonTailRequest,
@@ -149,7 +149,11 @@ export interface ServerCircuitProver {
    * Create a proof for the AVM circuit.
    * @param inputs - Inputs to the AVM circuit.
    */
-  getAvmProof(inputs: AvmCircuitInputs, signal?: AbortSignal, epochNumber?: number): Promise<ProofAndVerificationKey>;
+  getAvmProof(
+    inputs: AvmCircuitInputs,
+    signal?: AbortSignal,
+    epochNumber?: number,
+  ): Promise<AvmProofAndVerificationKey>;
 }
 
 /**

yarn-project/circuits.js/src/constants.gen.ts

@@ -202,6 +202,7 @@ export const RECURSIVE_PROOF_LENGTH = 439;
 export const NESTED_RECURSIVE_PROOF_LENGTH = 439;
 export const TUBE_PROOF_LENGTH = 439;
 export const VERIFICATION_KEY_LENGTH_IN_FIELDS = 128;
+export const AVM_VERIFICATION_KEY_LENGTH_IN_FIELDS = 66;
 export const SENDER_SELECTOR = 0;
 export const ADDRESS_SELECTOR = 1;
 export const STORAGE_ADDRESS_SELECTOR = 1;