@aztec/pxe 0.61.0 → 0.63.0

package/src/kernel_prover/kernel_prover.ts

@@ -23,7 +23,9 @@ import {
   VK_TREE_HEIGHT,
   VerificationKeyAsFields,
 } from '@aztec/circuits.js';
+import { hashVK } from '@aztec/circuits.js/hash';
 import { makeTuple } from '@aztec/foundation/array';
+import { vkAsFieldsMegaHonk } from '@aztec/foundation/crypto';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { assertLength } from '@aztec/foundation/serialize';
 import { pushTestData } from '@aztec/foundation/testing';
@@ -45,6 +47,7 @@ const NULL_PROVE_OUTPUT: PrivateKernelSimulateOutput<PrivateKernelCircuitPublicI
   outputWitness: new Map(),
   bytecode: Buffer.from([]),
 };
+
 /**
  * The KernelProver class is responsible for generating kernel proofs.
  * It takes a transaction request, its signature, and the simulation result as inputs, and outputs a proof
@@ -64,18 +67,30 @@ export class KernelProver {
    *
    * @param txRequest - The authenticated transaction request object.
    * @param executionResult - The execution result object containing nested executions and preimages.
+   * @param profile - Set true to profile the gate count for each circuit
+   * @param dryRun - Set true to skip the IVC proof generation (only simulation is run). Useful for profiling gate count without proof gen.
    * @returns A Promise that resolves to a KernelProverOutput object containing proof, public inputs, and output notes.
    * TODO(#7368) this should be refactored to not recreate the ACIR bytecode now that it operates on a program stack
    */
   async prove(
     txRequest: TxRequest,
     executionResult: PrivateExecutionResult,
+    profile: boolean = false,
+    dryRun: boolean = false,
   ): Promise<PrivateKernelSimulateOutput<PrivateKernelTailCircuitPublicInputs>> {
     const executionStack = [executionResult];
     let firstIteration = true;
 
     let output = NULL_PROVE_OUTPUT;
 
+    const gateCounts: { circuitName: string; gateCount: number }[] = [];
+    const addGateCount = async (circuitName: string, bytecode: Buffer) => {
+      const gateCount = (await this.proofCreator.computeGateCountForCircuit(bytecode, circuitName)) as number;
+      gateCounts.push({ circuitName, gateCount });
+
+      this.log.info(`Tx ${txRequest.hash()}: bb gates for ${circuitName} - ${gateCount}`);
+    };
+
     const noteHashLeafIndexMap = collectNoteHashLeafIndexMap(executionResult);
     const noteHashNullifierCounterMap = collectNoteHashNullifierCounterMap(executionResult);
     const enqueuedPublicFunctions = collectEnqueuedPublicFunctionCalls(executionResult);
@@ -100,6 +115,9 @@ export class KernelProver {
           // TODO(#7368) consider refactoring this redundant bytecode pushing
           acirs.push(output.bytecode);
           witnessStack.push(output.outputWitness);
+          if (profile) {
+            await addGateCount('private_kernel_reset', output.bytecode);
+          }
 
           resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
             output,
@@ -111,6 +129,7 @@ export class KernelProver {
       }
 
       const currentExecution = executionStack.pop()!;
+
       executionStack.push(...[...currentExecution.nestedExecutions].reverse());
 
       const functionName = await this.oracle.getDebugFunctionName(
@@ -118,13 +137,15 @@ export class KernelProver {
         currentExecution.publicInputs.callContext.functionSelector,
       );
 
-      const appVk = await this.proofCreator.computeAppCircuitVerificationKey(currentExecution.acir, functionName);
       // TODO(#7368): This used to be associated with getDebugFunctionName
       // TODO(#7368): Is there any way to use this with client IVC proving?
       acirs.push(currentExecution.acir);
       witnessStack.push(currentExecution.partialWitness);
+      if (profile) {
+        await addGateCount(functionName as string, currentExecution.acir);
+      }
 
-      const privateCallData = await this.createPrivateCallData(currentExecution, appVk.verificationKey);
+      const privateCallData = await this.createPrivateCallData(currentExecution);
 
       if (firstIteration) {
         const proofInput = new PrivateKernelInitCircuitPrivateInputs(
@@ -135,8 +156,12 @@ export class KernelProver {
         );
         pushTestData('private-kernel-inputs-init', proofInput);
         output = await this.proofCreator.simulateProofInit(proofInput);
+
         acirs.push(output.bytecode);
         witnessStack.push(output.outputWitness);
+        if (profile) {
+          await addGateCount('private_kernel_init', output.bytecode);
+        }
       } else {
         const previousVkMembershipWitness = await this.oracle.getVkMembershipWitness(output.verificationKey);
         const previousKernelData = new PrivateKernelData(
@@ -148,8 +173,12 @@ export class KernelProver {
         const proofInput = new PrivateKernelInnerCircuitPrivateInputs(previousKernelData, privateCallData);
         pushTestData('private-kernel-inputs-inner', proofInput);
         output = await this.proofCreator.simulateProofInner(proofInput);
+
         acirs.push(output.bytecode);
         witnessStack.push(output.outputWitness);
+        if (profile) {
+          await addGateCount('private_kernel_inner', output.bytecode);
+        }
       }
       firstIteration = false;
     }
@@ -164,8 +193,12 @@ export class KernelProver {
     while (resetBuilder.needsReset()) {
       const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
       output = await this.proofCreator.simulateProofReset(privateInputs);
+
       acirs.push(output.bytecode);
       witnessStack.push(output.outputWitness);
+      if (profile) {
+        await addGateCount('private_kernel_reset', output.bytecode);
+      }
 
       resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
         output,
@@ -192,18 +225,29 @@ export class KernelProver {
 
     pushTestData('private-kernel-inputs-ordering', privateInputs);
     const tailOutput = await this.proofCreator.simulateProofTail(privateInputs);
+
     acirs.push(tailOutput.bytecode);
     witnessStack.push(tailOutput.outputWitness);
+    if (profile) {
+      await addGateCount('private_kernel_tail', tailOutput.bytecode);
+      tailOutput.profileResult = { gateCounts };
+    }
 
     // TODO(#7368) how do we 'bincode' encode these inputs?
-    const ivcProof = await this.proofCreator.createClientIvcProof(acirs, witnessStack);
-    tailOutput.clientIvcProof = ivcProof;
+    if (!dryRun) {
+      const ivcProof = await this.proofCreator.createClientIvcProof(acirs, witnessStack);
+      tailOutput.clientIvcProof = ivcProof;
+    }
+
     return tailOutput;
   }
 
-  private async createPrivateCallData({ publicInputs }: PrivateExecutionResult, vk: VerificationKeyAsFields) {
+  private async createPrivateCallData({ publicInputs, vk: vkAsBuffer }: PrivateExecutionResult) {
     const { contractAddress, functionSelector } = publicInputs.callContext;
 
+    const vkAsFields = vkAsFieldsMegaHonk(vkAsBuffer);
+    const vk = new VerificationKeyAsFields(vkAsFields, hashVK(vkAsFields));
+
     const functionLeafMembershipWitness = await this.oracle.getFunctionMembershipWitness(
       contractAddress,
       functionSelector,

package/src/kernel_prover/test/test_circuit_prover.ts

@@ -107,6 +107,11 @@ export class TestPrivateKernelProver implements PrivateKernelProver {
     );
   }
 
+  public computeGateCountForCircuit(_bytecode: Buffer, _circuitName: string): Promise<number> {
+    // No gates in test prover
+    return Promise.resolve(0);
+  }
+
   computeAppCircuitVerificationKey(
     _bytecode: Buffer,
     _appCircuitName?: string | undefined,
@@ -117,10 +122,9 @@ export class TestPrivateKernelProver implements PrivateKernelProver {
     return Promise.resolve(appCircuitProofOutput);
   }
 
-  private makeEmptyKernelSimulateOutput<PublicInputsType>(
-    publicInputs: PublicInputsType,
-    circuitType: ProtocolArtifact,
-  ) {
+  private makeEmptyKernelSimulateOutput<
+    PublicInputsType extends PrivateKernelTailCircuitPublicInputs | PrivateKernelCircuitPublicInputs,
+  >(publicInputs: PublicInputsType, circuitType: ProtocolArtifact) {
     const kernelProofOutput: PrivateKernelSimulateOutput<PublicInputsType> = {
       publicInputs,
       verificationKey: ProtocolCircuitVks[circuitType].keyAsFields,

package/src/{note_processor/utils → note_decryption_utils}/add_public_values_to_payload.ts

@@ -1,7 +1,7 @@
 import { type L1NotePayload, Note } from '@aztec/circuit-types';
 import { ContractNotFoundError } from '@aztec/simulator';
 
-import { type PxeDatabase } from '../../database/pxe_database.js';
+import { type PxeDatabase } from '../database/pxe_database.js';
 
 /**
  * Merges privately and publicly delivered note values.

package/src/{note_processor/utils → note_decryption_utils}/produce_note_daos.ts

@@ -1,12 +1,11 @@
-import { type L1NotePayload, type PublicKey, type TxHash, type UnencryptedTxL2Logs } from '@aztec/circuit-types';
+import { type L1NotePayload, type PublicKey, type TxHash } from '@aztec/circuit-types';
 import { type Fr } from '@aztec/foundation/fields';
 import { type Logger } from '@aztec/foundation/log';
 import { type AcirSimulator } from '@aztec/simulator';
 
-import { type DeferredNoteDao } from '../../database/deferred_note_dao.js';
-import { IncomingNoteDao } from '../../database/incoming_note_dao.js';
-import { OutgoingNoteDao } from '../../database/outgoing_note_dao.js';
-import { type PxeDatabase } from '../../database/pxe_database.js';
+import { IncomingNoteDao } from '../database/incoming_note_dao.js';
+import { OutgoingNoteDao } from '../database/outgoing_note_dao.js';
+import { type PxeDatabase } from '../database/pxe_database.js';
 import { produceNoteDaosForKey } from './produce_note_daos_for_key.js';
 
 /**
@@ -17,7 +16,7 @@ import { produceNoteDaosForKey } from './produce_note_daos_for_key.js';
  *
  * @param simulator - An instance of AcirSimulator.
  * @param db - An instance of PxeDatabase.
- * @param ivpkM - The public counterpart to the secret key to be used in the decryption of incoming note logs.
+ * @param addressPoint - The public counterpart to the address secret, which is used in the decryption of incoming note logs.
  * @param ovpkM - The public counterpart to the secret key to be used in the decryption of outgoing note logs.
  * @param payload - An instance of l1NotePayload.
  * @param txHash - The hash of the transaction that created the note. Equivalent to the first nullifier of the transaction.
@@ -31,7 +30,7 @@ import { produceNoteDaosForKey } from './produce_note_daos_for_key.js';
 export async function produceNoteDaos(
   simulator: AcirSimulator,
   db: PxeDatabase,
-  ivpkM: PublicKey | undefined,
+  addressPoint: PublicKey | undefined,
   ovpkM: PublicKey | undefined,
   payload: L1NotePayload,
   txHash: TxHash,
@@ -39,34 +38,28 @@ export async function produceNoteDaos(
   dataStartIndexForTx: number,
   excludedIndices: Set<number>,
   logger: Logger,
-  unencryptedLogs: UnencryptedTxL2Logs,
 ): Promise<{
   incomingNote: IncomingNoteDao | undefined;
   outgoingNote: OutgoingNoteDao | undefined;
-  incomingDeferredNote: DeferredNoteDao | undefined;
-  outgoingDeferredNote: DeferredNoteDao | undefined;
 }> {
-  if (!ivpkM && !ovpkM) {
-    throw new Error('Both ivpkM and ovpkM are undefined. Cannot create note.');
+  if (!addressPoint && !ovpkM) {
+    throw new Error('Both addressPoint and ovpkM are undefined. Cannot create note.');
   }
 
   let incomingNote: IncomingNoteDao | undefined;
   let outgoingNote: OutgoingNoteDao | undefined;
-  let incomingDeferredNote: DeferredNoteDao | undefined;
-  let outgoingDeferredNote: DeferredNoteDao | undefined;
 
-  if (ivpkM) {
-    [incomingNote, incomingDeferredNote] = await produceNoteDaosForKey(
+  if (addressPoint) {
+    incomingNote = await produceNoteDaosForKey(
       simulator,
       db,
-      ivpkM,
+      addressPoint,
       payload,
       txHash,
       noteHashes,
       dataStartIndexForTx,
       excludedIndices,
       logger,
-      unencryptedLogs,
       IncomingNoteDao.fromPayloadAndNoteInfo,
     );
   }
@@ -87,7 +80,7 @@ export async function produceNoteDaos(
         ovpkM,
       );
     } else {
-      [outgoingNote, outgoingDeferredNote] = await produceNoteDaosForKey(
+      outgoingNote = await produceNoteDaosForKey(
         simulator,
         db,
         ovpkM,
@@ -97,7 +90,6 @@ export async function produceNoteDaos(
         dataStartIndexForTx,
         excludedIndices,
         logger,
-        unencryptedLogs,
         OutgoingNoteDao.fromPayloadAndNoteInfo,
       );
     }
@@ -106,7 +98,5 @@ export async function produceNoteDaos(
   return {
     incomingNote,
     outgoingNote,
-    incomingDeferredNote,
-    outgoingDeferredNote,
   };
 }

package/src/{note_processor/utils → note_decryption_utils}/produce_note_daos_for_key.ts

@@ -1,10 +1,9 @@
-import { type L1NotePayload, type Note, type TxHash, type UnencryptedTxL2Logs } from '@aztec/circuit-types';
+import { type L1NotePayload, type Note, type TxHash } from '@aztec/circuit-types';
 import { type Fr, type PublicKey } from '@aztec/circuits.js';
 import { type Logger } from '@aztec/foundation/log';
-import { type AcirSimulator, ContractNotFoundError } from '@aztec/simulator';
+import { type AcirSimulator } from '@aztec/simulator';
 
-import { DeferredNoteDao } from '../../database/deferred_note_dao.js';
-import { type PxeDatabase } from '../../database/pxe_database.js';
+import { type PxeDatabase } from '../database/pxe_database.js';
 import { getOrderedNoteItems } from './add_public_values_to_payload.js';
 import { type NoteInfo, bruteForceNoteInfo } from './brute_force_note_info.js';
 
@@ -18,7 +17,6 @@ export async function produceNoteDaosForKey<T>(
   dataStartIndexForTx: number,
   excludedIndices: Set<number>,
   logger: Logger,
-  unencryptedLogs: UnencryptedTxL2Logs,
   daoConstructor: (
     note: Note,
     payload: L1NotePayload,
@@ -26,9 +24,8 @@ export async function produceNoteDaosForKey<T>(
     dataStartIndexForTx: number,
     pkM: PublicKey,
   ) => T,
-): Promise<[T | undefined, DeferredNoteDao | undefined]> {
+): Promise<T | undefined> {
   let noteDao: T | undefined;
-  let deferredNoteDao: DeferredNoteDao | undefined;
 
   try {
     // We get the note by merging publicly and privately delivered note values.
@@ -49,14 +46,8 @@ export async function produceNoteDaosForKey<T>(
 
     noteDao = daoConstructor(note, payload, noteInfo, dataStartIndexForTx, pkM);
   } catch (e) {
-    if (e instanceof ContractNotFoundError) {
-      logger.warn(e.message);
-
-      deferredNoteDao = new DeferredNoteDao(pkM, payload, txHash, noteHashes, dataStartIndexForTx, unencryptedLogs);
-    } else {
-      logger.error(`Could not process note because of "${e}". Discarding note...`);
-    }
+    logger.error(`Could not process note because of "${e}". Discarding note...`);
   }
 
-  return [noteDao, deferredNoteDao];
+  return noteDao;
 }

package/src/pxe_http/pxe_http_server.ts

@@ -1,38 +1,5 @@
-import {
-  AuthWitness,
-  CompleteAddress,
-  CountedNoteLog,
-  CountedPublicExecutionRequest,
-  EncryptedL2Log,
-  EncryptedL2NoteLog,
-  EncryptedNoteL2BlockL2Logs,
-  ExtendedNote,
-  ExtendedUnencryptedL2Log,
-  L2Block,
-  LogId,
-  Note,
-  NullifierMembershipWitness,
-  type PXE,
-  PrivateExecutionResult,
-  SiblingPath,
-  Tx,
-  TxEffect,
-  TxExecutionRequest,
-  TxHash,
-  TxProvingResult,
-  TxReceipt,
-  TxSimulationResult,
-  UnencryptedL2BlockL2Logs,
-  UnencryptedL2Log,
-  UniqueNote,
-} from '@aztec/circuit-types';
-import { FunctionSelector, PrivateCircuitPublicInputs, PublicKeys } from '@aztec/circuits.js';
-import { NoteSelector } from '@aztec/foundation/abi';
-import { AztecAddress } from '@aztec/foundation/aztec-address';
-import { Buffer32 } from '@aztec/foundation/buffer';
-import { EthAddress } from '@aztec/foundation/eth-address';
-import { Fr, GrumpkinScalar, Point } from '@aztec/foundation/fields';
-import { JsonRpcServer, createNamespacedJsonRpcServer } from '@aztec/foundation/json-rpc/server';
+import { type PXE, PXESchema } from '@aztec/circuit-types';
+import { createNamespacedSafeJsonRpcServer, createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
 
 import http from 'http';
 
@@ -40,50 +7,8 @@ import http from 'http';
  * Wraps an instance of Private eXecution Environment (PXE) implementation to a JSON RPC HTTP interface.
  * @returns A new instance of the HTTP server.
  */
-export function createPXERpcServer(pxeService: PXE): JsonRpcServer {
-  return new JsonRpcServer(
-    pxeService,
-    {
-      CompleteAddress,
-      AztecAddress,
-      TxExecutionRequest,
-      ExtendedUnencryptedL2Log,
-      FunctionSelector,
-      TxHash,
-      Buffer32,
-      EthAddress,
-      Point,
-      Fr,
-      GrumpkinScalar,
-      Note,
-      ExtendedNote,
-      PublicKeys,
-      UniqueNote,
-      SiblingPath,
-      AuthWitness,
-      L2Block,
-      TxEffect,
-      LogId,
-    },
-    {
-      EncryptedNoteL2BlockL2Logs,
-      EncryptedL2NoteLog,
-      EncryptedL2Log,
-      UnencryptedL2Log,
-      NoteSelector,
-      NullifierMembershipWitness,
-      TxSimulationResult,
-      TxProvingResult,
-      PrivateCircuitPublicInputs,
-      PrivateExecutionResult,
-      CountedPublicExecutionRequest,
-      CountedNoteLog,
-      Tx,
-      TxReceipt,
-      UnencryptedL2BlockL2Logs,
-    },
-    ['start', 'stop'],
-  );
+export function createPXERpcServer(pxeService: PXE) {
+  return createSafeJsonRpcServer(pxeService, PXESchema);
 }
 
 /**
@@ -93,8 +18,7 @@ export function createPXERpcServer(pxeService: PXE): JsonRpcServer {
  * @returns A running http server.
  */
 export function startPXEHttpServer(pxeService: PXE, port: string | number): http.Server {
-  const pxeServer = createPXERpcServer(pxeService);
-  const rpcServer = createNamespacedJsonRpcServer([{ pxe: pxeServer }]);
+  const rpcServer = createNamespacedSafeJsonRpcServer({ pxe: [pxeService, PXESchema] });
 
   const app = rpcServer.getApp();
   const httpServer = http.createServer(app.callback());

package/src/pxe_service/create_pxe_service.ts

@@ -31,11 +31,17 @@ export async function createPXEService(
     typeof useLogSuffix === 'boolean' ? (useLogSuffix ? randomBytes(3).toString('hex') : undefined) : useLogSuffix;
 
   const l1Contracts = await aztecNode.getL1ContractAddresses();
-  const storeConfig = { dataDirectory: config.dataDirectory, l1Contracts };
+  const configWithContracts = {
+    ...config,
+    l1Contracts,
+  } as PXEServiceConfig;
+
   const keyStore = new KeyStore(
-    await createStore('pxe_key_store', storeConfig, createDebugLogger('aztec:pxe:keystore:lmdb')),
+    await createStore('pxe_key_store', configWithContracts, createDebugLogger('aztec:pxe:keystore:lmdb')),
+  );
+  const db = new KVPxeDatabase(
+    await createStore('pxe_data', configWithContracts, createDebugLogger('aztec:pxe:data:lmdb')),
   );
-  const db = new KVPxeDatabase(await createStore('pxe_data', storeConfig, createDebugLogger('aztec:pxe:data:lmdb')));
 
   const prover = proofCreator ?? (await createProver(config, logSuffix));
   const server = new PXEService(keyStore, aztecNode, db, prover, config, logSuffix);

package/src/pxe_service/error_enriching.ts

@@ -1,7 +1,7 @@
 import { type SimulationError, isNoirCallStackUnresolved } from '@aztec/circuit-types';
 import { AztecAddress, Fr, FunctionSelector, PUBLIC_DISPATCH_SELECTOR } from '@aztec/circuits.js';
 import { type DebugLogger } from '@aztec/foundation/log';
-import { resolveAssertionMessage, resolveOpcodeLocations } from '@aztec/simulator';
+import { resolveAssertionMessageFromRevertData, resolveOpcodeLocations } from '@aztec/simulator';
 
 import { type ContractDataOracle, type PxeDatabase } from '../index.js';
 
@@ -64,17 +64,24 @@ export async function enrichPublicSimulationError(
   // To be able to resolve the assertion message, we need to use the information from the public dispatch function,
   // no matter what the call stack selector points to (since we've modified it to point to the target function).
   // We should remove this because the AVM (or public protocol) shouldn't be aware of the public dispatch calling convention.
+
+  const artifact = await contractDataOracle.getFunctionArtifact(
+    originalFailingFunction.contractAddress,
+    FunctionSelector.fromField(new Fr(PUBLIC_DISPATCH_SELECTOR)),
+  );
+  const assertionMessage = resolveAssertionMessageFromRevertData(err.revertData, artifact);
+  if (assertionMessage) {
+    err.setOriginalMessage(err.getOriginalMessage() + `${assertionMessage}`);
+  }
+
   const debugInfo = await contractDataOracle.getFunctionDebugMetadata(
     originalFailingFunction.contractAddress,
     FunctionSelector.fromField(new Fr(PUBLIC_DISPATCH_SELECTOR)),
   );
+
   const noirCallStack = err.getNoirCallStack();
   if (debugInfo) {
     if (isNoirCallStackUnresolved(noirCallStack)) {
-      const assertionMessage = resolveAssertionMessage(noirCallStack, debugInfo);
-      if (assertionMessage) {
-        err.setOriginalMessage(err.getOriginalMessage() + `: ${assertionMessage}`);
-      }
       try {
         // Public functions are simulated as a single Brillig entry point.
         // Thus, we can safely assume here that the Brillig function id is `0`.