chore: granular public simulation benchmarks (#6924)

Allow the PXE to send classes that get registered to its node to allow
better debug logging

Consolidate the kernel phases enum
Add benchmarking stats for AVM simulation
Add benchmarking stats for public DB accesses made by avm simulation

yarn-project/archiver/package.json

@@ -67,6 +67,7 @@
     "ws": "^8.13.0"
   },
   "devDependencies": {
+    "@aztec/noir-contracts.js": "workspace:^",
     "@jest/globals": "^29.5.0",
     "@types/debug": "^4.1.7",
     "@types/jest": "^29.5.0",

yarn-project/archiver/src/archiver/archiver.ts

@@ -22,6 +22,7 @@ import {
   isValidUnconstrainedFunctionMembershipProof,
 } from '@aztec/circuits.js/contract';
 import { createEthereumChain } from '@aztec/ethereum';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { type AztecAddress } from '@aztec/foundation/aztec-address';
 import { type EthAddress } from '@aztec/foundation/eth-address';
 import { Fr } from '@aztec/foundation/fields';
@@ -489,4 +490,12 @@ export class Archiver implements ArchiveSource {
   getContractClassIds(): Promise<Fr[]> {
     return this.store.getContractClassIds();
   }
+
+  addContractArtifact(address: AztecAddress, artifact: ContractArtifact): Promise<void> {
+    return this.store.addContractArtifact(address, artifact);
+  }
+
+  getContractArtifact(address: AztecAddress): Promise<ContractArtifact | undefined> {
+    return this.store.getContractArtifact(address);
+  }
 }

yarn-project/archiver/src/archiver/archiver_store.ts

@@ -15,6 +15,7 @@ import {
   type UnencryptedL2BlockL2Logs,
 } from '@aztec/circuit-types';
 import { type Fr } from '@aztec/circuits.js';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { type AztecAddress } from '@aztec/foundation/aztec-address';
 import {
   type ContractClassPublic,
@@ -191,4 +192,7 @@ export interface ArchiverDataStore {
 
   /** Returns the list of all class ids known by the archiver. */
   getContractClassIds(): Promise<Fr[]>;
+
+  addContractArtifact(address: AztecAddress, contract: ContractArtifact): Promise<void>;
+  getContractArtifact(address: AztecAddress): Promise<ContractArtifact | undefined>;
 }

yarn-project/archiver/src/archiver/kv_archiver_store/contract_artifacts_store.test.ts

@@ -0,0 +1,22 @@
+import { AztecAddress } from '@aztec/circuits.js';
+import { openTmpStore } from '@aztec/kv-store/utils';
+import { BenchmarkingContractArtifact } from '@aztec/noir-contracts.js/Benchmarking';
+
+import { beforeEach } from '@jest/globals';
+
+import { KVArchiverDataStore } from './kv_archiver_store.js';
+
+describe('Contract Artifact Store', () => {
+  let archiverStore: KVArchiverDataStore;
+
+  beforeEach(() => {
+    archiverStore = new KVArchiverDataStore(openTmpStore());
+  });
+
+  it('Should add and return contract artifacts', async () => {
+    const artifact = BenchmarkingContractArtifact;
+    const address = AztecAddress.random();
+    await archiverStore.addContractArtifact(address, artifact);
+    await expect(archiverStore.getContractArtifact(address)).resolves.toEqual(artifact);
+  });
+});

yarn-project/archiver/src/archiver/kv_archiver_store/contract_artifacts_store.ts

@@ -0,0 +1,22 @@
+import { type AztecAddress } from '@aztec/circuits.js';
+import { type ContractArtifact } from '@aztec/foundation/abi';
+import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
+import { contractArtifactFromBuffer, contractArtifactToBuffer } from '@aztec/types/abi';
+
+export class ContractArtifactsStore {
+  #contractArtifacts: AztecMap<string, Buffer>;
+
+  constructor(db: AztecKVStore) {
+    this.#contractArtifacts = db.openMap('archiver_contract_artifacts');
+  }
+
+  addContractArtifact(address: AztecAddress, contractArtifact: ContractArtifact): Promise<void> {
+    return this.#contractArtifacts.set(address.toString(), contractArtifactToBuffer(contractArtifact));
+  }
+
+  getContractArtifact(address: AztecAddress): ContractArtifact | undefined {
+    const contractArtifact = this.#contractArtifacts.get(address.toString());
+    // TODO(@spalladino): AztecMap lies and returns Uint8Arrays instead of Buffers, hence the extra Buffer.from.
+    return contractArtifact && contractArtifactFromBuffer(Buffer.from(contractArtifact));
+  }
+}

yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.ts

@@ -15,6 +15,7 @@ import {
   type UnencryptedL2BlockL2Logs,
 } from '@aztec/circuit-types';
 import { type Fr } from '@aztec/circuits.js';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { type AztecAddress } from '@aztec/foundation/aztec-address';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { type AztecKVStore } from '@aztec/kv-store';
@@ -29,6 +30,7 @@ import { type ArchiverDataStore, type ArchiverL1SynchPoint } from '../archiver_s
 import { type DataRetrieval } from '../data_retrieval.js';
 import { BlockBodyStore } from './block_body_store.js';
 import { BlockStore } from './block_store.js';
+import { ContractArtifactsStore } from './contract_artifacts_store.js';
 import { ContractClassStore } from './contract_class_store.js';
 import { ContractInstanceStore } from './contract_instance_store.js';
 import { LogStore } from './log_store.js';
@@ -44,6 +46,7 @@ export class KVArchiverDataStore implements ArchiverDataStore {
   #messageStore: MessageStore;
   #contractClassStore: ContractClassStore;
   #contractInstanceStore: ContractInstanceStore;
+  #contractArtifactStore: ContractArtifactsStore;
 
   #log = createDebugLogger('aztec:archiver:data-store');
 
@@ -54,6 +57,15 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     this.#messageStore = new MessageStore(db);
     this.#contractClassStore = new ContractClassStore(db);
     this.#contractInstanceStore = new ContractInstanceStore(db);
+    this.#contractArtifactStore = new ContractArtifactsStore(db);
+  }
+
+  getContractArtifact(address: AztecAddress): Promise<ContractArtifact | undefined> {
+    return Promise.resolve(this.#contractArtifactStore.getContractArtifact(address));
+  }
+
+  addContractArtifact(address: AztecAddress, contract: ContractArtifact): Promise<void> {
+    return this.#contractArtifactStore.addContractArtifact(address, contract);
   }
 
   getContractClass(id: Fr): Promise<ContractClassPublic | undefined> {

yarn-project/archiver/src/archiver/memory_archiver_store/memory_archiver_store.ts

@@ -17,6 +17,7 @@ import {
   type UnencryptedL2BlockL2Logs,
 } from '@aztec/circuit-types';
 import { Fr, INITIAL_L2_BLOCK_NUM } from '@aztec/circuits.js';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { type AztecAddress } from '@aztec/foundation/aztec-address';
 import {
   type ContractClassPublic,
@@ -71,6 +72,8 @@ export class MemoryArchiverStore implements ArchiverDataStore {
    */
   private l1ToL2Messages = new L1ToL2MessageStore();
 
+  private contractArtifacts: Map<string, ContractArtifact> = new Map();
+
   private contractClasses: Map<string, ContractClassPublic> = new Map();
 
   private privateFunctions: Map<string, ExecutablePrivateFunctionWithMembershipProof[]> = new Map();
@@ -438,4 +441,13 @@ export class MemoryArchiverStore implements ArchiverDataStore {
       messagesSynchedTo: this.lastL1BlockNewMessages,
     });
   }
+
+  public addContractArtifact(address: AztecAddress, contract: ContractArtifact): Promise<void> {
+    this.contractArtifacts.set(address.toString(), contract);
+    return Promise.resolve();
+  }
+
+  public getContractArtifact(address: AztecAddress): Promise<ContractArtifact | undefined> {
+    return Promise.resolve(this.contractArtifacts.get(address.toString()));
+  }
 }

yarn-project/archiver/tsconfig.json

@@ -29,6 +29,9 @@
     },
     {
       "path": "../types"
+    },
+    {
+      "path": "../noir-contracts.js"
     }
   ],
   "include": ["src"]

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -45,6 +45,7 @@ import {
 } from '@aztec/circuits.js';
 import { computePublicDataTreeLeafSlot } from '@aztec/circuits.js/hash';
 import { type L1ContractAddresses, createEthereumChain } from '@aztec/ethereum';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { type AztecKVStore } from '@aztec/kv-store';
@@ -764,6 +765,10 @@ export class AztecNodeService implements AztecNode {
     });
   }
 
+  public addContractArtifact(address: AztecAddress, artifact: ContractArtifact): Promise<void> {
+    return this.contractDataSource.addContractArtifact(address, artifact);
+  }
+
   /**
    * Returns an instance of MerkleTreeOperations having first ensured the world state is fully synched
    * @param blockNumber - The block number at which to get the data.

yarn-project/aztec-node/tsconfig.json

@@ -36,6 +36,9 @@
     {
       "path": "../p2p"
     },
+    {
+      "path": "../protocol-contracts"
+    },
     {
       "path": "../prover-client"
     },
@@ -50,9 +53,6 @@
     },
     {
       "path": "../world-state"
-    },
-    {
-      "path": "../protocol-contracts"
     }
   ],
   "include": ["src"]

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

@@ -4,7 +4,6 @@ import {
   type PublicInputsAndRecursiveProof,
   type PublicKernelNonTailRequest,
   type PublicKernelTailRequest,
-  PublicKernelType,
   type ServerCircuitProver,
   makePublicInputsAndRecursiveProof,
 } from '@aztec/circuit-types';
@@ -181,7 +180,7 @@ export class BBNativeRollupProver implements ServerCircuitProver {
   ): Promise<PublicInputsAndRecursiveProof<PublicKernelCircuitPublicInputs>> {
     const kernelOps = PublicKernelArtifactMapping[kernelRequest.type];
     if (kernelOps === undefined) {
-      throw new Error(`Unable to prove kernel type ${PublicKernelType[kernelRequest.type]}`);
+      throw new Error(`Unable to prove kernel type ${kernelRequest.type}`);
     }
 
     // We may need to convert the recursive proof into fields format

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

@@ -3,7 +3,6 @@ import {
   type PublicInputsAndRecursiveProof,
   type PublicKernelNonTailRequest,
   type PublicKernelTailRequest,
-  PublicKernelType,
   type ServerCircuitProver,
   makePublicInputsAndRecursiveProof,
 } from '@aztec/circuit-types';
@@ -265,7 +264,7 @@ export class TestCircuitProver implements ServerCircuitProver {
     const timer = new Timer();
     const kernelOps = SimulatedPublicKernelArtifactMapping[kernelRequest.type];
     if (kernelOps === undefined) {
-      throw new Error(`Unable to prove for kernel type ${PublicKernelType[kernelRequest.type]}`);
+      throw new Error(`Unable to prove for kernel type ${kernelRequest.type}`);
     }
     const witnessMap = kernelOps.convertInputs(kernelRequest.inputs);
 

yarn-project/circuit-types/src/index.ts

@@ -1,22 +1,22 @@
+export { CompleteAddress, GrumpkinPrivateKey, type PartialAddress, type PublicKey } from '@aztec/circuits.js';
+export * from './auth_witness.js';
+export * from './aztec_node/rpc/index.js';
+export * from './body.js';
 export * from './function_call.js';
-export * from './notes/index.js';
-export * from './messaging/index.js';
+export * from './interfaces/index.js';
 export * from './l2_block.js';
-export * from './body.js';
 export * from './l2_block_downloader/index.js';
 export * from './l2_block_source.js';
-export * from './public_data_witness.js';
-export * from './tx_effect.js';
 export * from './logs/index.js';
 export * from './merkle_tree_id.js';
+export * from './messaging/index.js';
 export * from './mocks.js';
+export * from './notes/index.js';
+export * from './packed_values.js';
+export * from './public_data_witness.js';
 export * from './public_data_write.js';
-export * from './simulation_error.js';
 export * from './sibling_path/index.js';
+export * from './simulation_error.js';
 export * from './tx/index.js';
+export * from './tx_effect.js';
 export * from './tx_execution_request.js';
-export * from './packed_values.js';
-export * from './interfaces/index.js';
-export * from './auth_witness.js';
-export * from './aztec_node/rpc/index.js';
-export { CompleteAddress, type PublicKey, type PartialAddress, GrumpkinPrivateKey } from '@aztec/circuits.js';

yarn-project/circuit-types/src/interfaces/aztec-node.ts

@@ -7,6 +7,7 @@ import {
   type PUBLIC_DATA_TREE_HEIGHT,
 } from '@aztec/circuits.js';
 import { type L1ContractAddresses } from '@aztec/ethereum';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { type AztecAddress } from '@aztec/foundation/aztec-address';
 import { type Fr } from '@aztec/foundation/fields';
 import {
@@ -218,6 +219,13 @@ export interface AztecNode {
    */
   getProtocolContractAddresses(): Promise<ProtocolContractAddresses>;
 
+  /**
+   * Method to add a contract artifact to the database.
+   * @param aztecAddress
+   * @param artifact
+   */
+  addContractArtifact(address: AztecAddress, artifact: ContractArtifact): Promise<void>;
+
   /**
    * Gets up to `limit` amount of logs starting from `from`.
    * @param from - Number of the L2 block to which corresponds the first logs to be returned.

yarn-project/circuit-types/src/stats/metrics.ts

@@ -25,6 +25,18 @@ export interface Metric {
 
 /** Metric definitions to track from benchmarks. */
 export const Metrics = [
+  {
+    name: 'public_db_access_time_ms',
+    groupBy: 'chain-length',
+    description: 'Time to access a database.',
+    events: ['public-db-access'],
+  },
+  {
+    name: 'avm_simulation_time_ms',
+    groupBy: 'app-circuit-name',
+    description: 'Time to simulate an AVM circuit.',
+    events: ['avm-simulation'],
+  },
   {
     name: 'proof_construction_time_sha256_ms',
     groupBy: 'threads',

yarn-project/circuit-types/src/stats/stats.ts

@@ -99,6 +99,21 @@ export type CircuitSimulationStats = {
   outputSize: number;
 };
 
+export type PublicDBAccessStats = {
+  eventName: 'public-db-access';
+  duration: number;
+  operation: string;
+};
+
+export type AvmSimulationStats = {
+  /** name of the event. */
+  eventName: 'avm-simulation';
+  /** Name of the circuit. */
+  appCircuitName: string;
+  /** Duration in ms. */
+  duration: number;
+};
+
 /** Stats for witness generation. */
 export type CircuitWitnessGenerationStats = {
   /** name of the event. */
@@ -247,17 +262,19 @@ export type TxAddedToPoolStats = {
 
 /** Stats emitted in structured logs with an `eventName` for tracking. */
 export type Stats =
-  | ProofConstructed
-  | L1PublishStats
-  | NodeSyncedChainHistoryStats
-  | CircuitSimulationStats
+  | AvmSimulationStats
   | CircuitProvingStats
+  | CircuitSimulationStats
   | CircuitWitnessGenerationStats
+  | PublicDBAccessStats
+  | L1PublishStats
   | L2BlockBuiltStats
   | L2BlockHandledStats
+  | NodeSyncedChainHistoryStats
   | NoteProcessorCaughtUpStats
-  | TxAddedToPoolStats
-  | TreeInsertionStats;
+  | ProofConstructed
+  | TreeInsertionStats
+  | TxAddedToPoolStats;
 
 /** Set of event names across emitted stats. */
 export type StatsEventName = Stats['eventName'];

yarn-project/circuit-types/src/tx/index.ts

@@ -1,8 +1,8 @@
-export * from './tx.js';
+export * from './processed_tx.js';
+export * from './public_simulation_output.js';
 export * from './simulated_tx.js';
+export * from './tx.js';
 export * from './tx_hash.js';
 export * from './tx_receipt.js';
-export * from './processed_tx.js';
-export * from './public_simulation_output.js';
-export * from './validator/tx_validator.js';
 export * from './validator/aggregate_tx_validator.js';
+export * from './validator/tx_validator.js';

yarn-project/circuit-types/src/tx/processed_tx.ts

@@ -31,11 +31,11 @@ import {
  * Used to communicate to the prover which type of circuit to prove
  */
 export enum PublicKernelType {
-  NON_PUBLIC,
-  SETUP,
-  APP_LOGIC,
-  TEARDOWN,
-  TAIL,
+  NON_PUBLIC = 'non-public',
+  SETUP = 'setup',
+  APP_LOGIC = 'app-logic',
+  TEARDOWN = 'teardown',
+  TAIL = 'tail',
 }
 
 export type PublicKernelTailRequest = {