@aztec/pxe 0.0.0-test.1 → 0.0.1-commit.03f7ef2

package/src/contract_function_simulator/proxied_contract_data_source.ts

@@ -0,0 +1,83 @@
+import { FunctionSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractOverrides } from '@aztec/stdlib/tx';
+
+import type { ContractDataProvider } from '../storage/contract_data_provider/contract_data_provider.js';
+
+/*
+ * Proxy generator for a ContractDataProvider that allows overriding contract instances and artifacts, so
+ * the contract function simulator can execute different bytecode on certain addresses. An example use case
+ * would be overriding your own account contract so that valid signatures don't have to be provided while simulating.
+ */
+export class ProxiedContractDataProviderFactory {
+  static create(contractDataProvider: ContractDataProvider, overrides?: ContractOverrides) {
+    if (!overrides) {
+      return contractDataProvider;
+    }
+
+    return new Proxy(contractDataProvider, {
+      get(target, prop: keyof ContractDataProvider) {
+        switch (prop) {
+          case 'getContractInstance': {
+            return async (address: AztecAddress) => {
+              if (overrides[address.toString()]) {
+                const { instance } = overrides[address.toString()]!;
+                instance.address = address;
+                const realInstance = await target.getContractInstance(address);
+                if (!realInstance) {
+                  throw new Error(`Contract instance not found for address: ${address}`);
+                }
+                instance.currentContractClassId = realInstance.currentContractClassId;
+                instance.originalContractClassId = realInstance.originalContractClassId;
+                return instance;
+              } else {
+                return target.getContractInstance(address);
+              }
+            };
+          }
+          case 'getFunctionArtifact': {
+            return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
+              if (overrides[contractAddress.toString()]) {
+                const { artifact } = overrides[contractAddress.toString()]!;
+                const functions = artifact.functions;
+                for (let i = 0; i < functions.length; i++) {
+                  const fn = functions[i];
+                  const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+                  if (fnSelector.equals(selector)) {
+                    return fn;
+                  }
+                }
+              } else {
+                return target.getFunctionArtifact(contractAddress, selector);
+              }
+            };
+          }
+          case 'getFunctionArtifactWithDebugMetadata': {
+            return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
+              if (overrides[contractAddress.toString()]) {
+                const { artifact } = overrides[contractAddress.toString()]!;
+                const functions = artifact.functions;
+                for (let i = 0; i < functions.length; i++) {
+                  const fn = functions[i];
+                  const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+                  if (fnSelector.equals(selector)) {
+                    return fn;
+                  }
+                }
+              } else {
+                return target.getFunctionArtifactWithDebugMetadata(contractAddress, selector);
+              }
+            };
+          }
+          default: {
+            const value = Reflect.get(target, prop);
+            if (typeof value === 'function') {
+              return value.bind(target);
+            }
+            return value;
+          }
+        }
+      },
+    });
+  }
+}

package/src/contract_function_simulator/proxied_node.ts

@@ -0,0 +1,33 @@
+import { Timer } from '@aztec/foundation/timer';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { NodeStats } from '@aztec/stdlib/tx';
+
+/*
+ * Proxy generator for an AztecNode that tracks the time taken for each RPC call.
+ */
+export type ProxiedNode = AztecNode & { getStats(): NodeStats };
+
+export class ProxiedNodeFactory {
+  static create(node: AztecNode) {
+    const stats: Partial<Record<keyof AztecNode, { times: number[] }>> = {};
+    return new Proxy(node, {
+      get(target, prop: keyof ProxiedNode) {
+        if (prop === 'getStats') {
+          return () => {
+            return stats;
+          };
+        } else {
+          return async function (...args: any[]) {
+            if (!stats[prop]) {
+              stats[prop] = { times: [] };
+            }
+            const timer = new Timer();
+            const result = await (target[prop] as any).apply(target, args);
+            stats[prop].times.push(timer.ms());
+            return result;
+          };
+        }
+      },
+    }) as ProxiedNode;
+  }
+}

package/src/debug/pxe_debug_utils.ts

@@ -0,0 +1,48 @@
+import type { NoteDao, NotesFilter } from '@aztec/stdlib/note';
+
+import type { PXE } from '../pxe.js';
+import type { ContractDataProvider } from '../storage/contract_data_provider/contract_data_provider.js';
+import type { NoteDataProvider } from '../storage/note_data_provider/note_data_provider.js';
+
+/**
+ * Methods provided by this class might help debugging but must not be used in production.
+ * No backwards compatibility or API stability should be expected. Use at your own risk.
+ */
+export class PXEDebugUtils {
+  #pxe: PXE | undefined = undefined;
+
+  constructor(
+    private contractDataProvider: ContractDataProvider,
+    private noteDataProvider: NoteDataProvider,
+  ) {}
+
+  /**
+   * Not injected through constructor since they're are co-dependant.
+   */
+  public setPXE(pxe: PXE) {
+    this.#pxe = pxe;
+  }
+
+  /**
+   * A debugging utility to get notes based on the provided filter.
+   *
+   * Note that this should not be used in production code because the structure of notes is considered to be
+   * an implementation detail of contracts. This is only meant to be used for debugging purposes. If you need to obtain
+   * note-related information in production code, please implement a custom utility function on your contract and call
+   * that function instead (e.g. `get_balance(owner: AztecAddress) -> u128` utility function on a Token contract).
+   *
+   * @param filter - The filter to apply to the notes.
+   * @returns The requested notes.
+   */
+  public async getNotes(filter: NotesFilter): Promise<NoteDao[]> {
+    if (!this.#pxe) {
+      throw new Error('Cannot getNotes because no PXE is set');
+    }
+
+    // We need to manually trigger private state sync to have a guarantee that all the notes are available.
+    const call = await this.contractDataProvider.getFunctionCall('sync_private_state', [], filter.contractAddress);
+    await this.#pxe.simulateUtility(call);
+
+    return this.noteDataProvider.getNotes(filter);
+  }
+}

package/src/entrypoints/client/bundle/index.ts

@@ -1,5 +1,6 @@
-export * from '../../../pxe_service/index.js';
+export * from '../../../pxe.js';
 export * from '../../../config/index.js';
+export * from '../../../error_enriching.js';
 export * from '../../../storage/index.js';
 export * from './utils.js';
-export { PXEOracleInterface } from '../../../pxe_oracle_interface/index.js';
+export type { PXECreationOptions } from '../../pxe_creation_options.js';

package/src/entrypoints/client/bundle/utils.ts

@@ -1,58 +1,58 @@
-import { BBWASMBundlePrivateKernelProver } from '@aztec/bb-prover/wasm/bundle';
+import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
+import { randomBytes } from '@aztec/foundation/crypto/random';
 import { createLogger } from '@aztec/foundation/log';
 import { createStore } from '@aztec/kv-store/indexeddb';
 import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
 import { WASMSimulator } from '@aztec/simulator/client';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
-import type { PXEServiceConfig } from '../../../config/index.js';
-import { PXEService } from '../../../pxe_service/pxe_service.js';
-import type { PXECreationOptions } from '../pxe_creation_options.js';
+import type { PXEConfig } from '../../../config/index.js';
+import { PXE } from '../../../pxe.js';
+import type { PXECreationOptions } from '../../pxe_creation_options.js';
 
 /**
- * Create and start an PXEService instance with the given AztecNode.
+ * Create and start an PXE instance with the given AztecNode.
  * If no keyStore or database is provided, it will use KeyStore and MemoryDB as default values.
- * Returns a Promise that resolves to the started PXEService instance.
+ * Returns a Promise that resolves to the started PXE instance.
  *
  * @param aztecNode - The AztecNode instance to be used by the server.
- * @param config - The PXE Service Config to use
- * @param options - (Optional) Optional information for creating an PXEService.
- * @returns A Promise that resolves to the started PXEService instance.
+ * @param config - The PXE Config to use
+ * @param options - (Optional) Optional information for creating an PXE.
+ * @returns A Promise that resolves to the started PXE instance.
  */
-export async function createPXEService(
+export async function createPXE(
   aztecNode: AztecNode,
-  config: PXEServiceConfig,
+  config: PXEConfig,
   options: PXECreationOptions = { loggers: {} },
 ) {
+  const logSuffix =
+    typeof options.useLogSuffix === 'boolean'
+      ? options.useLogSuffix
+        ? randomBytes(3).toString('hex')
+        : undefined
+      : options.useLogSuffix;
+
+  const loggers = options.loggers ?? {};
+
   const l1Contracts = await aztecNode.getL1ContractAddresses();
   const configWithContracts = {
     ...config,
     l1Contracts,
-  } as PXEServiceConfig;
-
-  const store = await createStore(
-    'pxe_data',
-    configWithContracts,
-    options.loggers.store ?? createLogger('pxe:data:indexeddb'),
-  );
-
-  const simulationProvider = new WASMSimulator();
-  const prover =
-    options.prover ??
-    new BBWASMBundlePrivateKernelProver(
-      simulationProvider,
-      16,
-      options.loggers.prover ?? createLogger('bb:wasm:bundle'),
-    );
+  } as PXEConfig;
+
+  const storeLogger = loggers.store ? loggers.store : createLogger('pxe:data:idb' + (logSuffix ? `:${logSuffix}` : ''));
+
+  const store = options.store ?? (await createStore('pxe_data', configWithContracts, storeLogger));
+
+  const simulator = options.simulator ?? new WASMSimulator();
+  const proverLogger = loggers.prover
+    ? loggers.prover
+    : createLogger('pxe:bb:wasm:bundle' + (logSuffix ? `:${logSuffix}` : ''));
+
+  const prover = options.prover ?? new BBBundlePrivateKernelProver(simulator, proverLogger);
   const protocolContractsProvider = new BundledProtocolContractsProvider();
-  const pxe = await PXEService.create(
-    aztecNode,
-    store,
-    prover,
-    simulationProvider,
-    protocolContractsProvider,
-    config,
-    options.loggers.pxe ?? createLogger('pxe:service'),
-  );
+
+  const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));
+  const pxe = await PXE.create(aztecNode, store, prover, simulator, protocolContractsProvider, config, pxeLogger);
   return pxe;
 }

package/src/entrypoints/client/lazy/index.ts

@@ -1,5 +1,6 @@
-export * from '../../../pxe_service/index.js';
+export * from '../../../pxe.js';
 export * from '../../../config/index.js';
 export * from '../../../storage/index.js';
+export * from '../../../error_enriching.js';
 export * from './utils.js';
-export { PXEOracleInterface } from '../../../pxe_oracle_interface/index.js';
+export { type PXECreationOptions } from '../../pxe_creation_options.js';

package/src/entrypoints/client/lazy/utils.ts

@@ -1,53 +1,58 @@
-import { BBWASMLazyPrivateKernelProver } from '@aztec/bb-prover/wasm/lazy';
+import { BBLazyPrivateKernelProver } from '@aztec/bb-prover/client/lazy';
+import { randomBytes } from '@aztec/foundation/crypto/random';
 import { createLogger } from '@aztec/foundation/log';
 import { createStore } from '@aztec/kv-store/indexeddb';
 import { LazyProtocolContractsProvider } from '@aztec/protocol-contracts/providers/lazy';
 import { WASMSimulator } from '@aztec/simulator/client';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
-import type { PXEServiceConfig } from '../../../config/index.js';
-import { PXEService } from '../../../pxe_service/pxe_service.js';
-import type { PXECreationOptions } from '../pxe_creation_options.js';
+import type { PXEConfig } from '../../../config/index.js';
+import { PXE } from '../../../pxe.js';
+import type { PXECreationOptions } from '../../pxe_creation_options.js';
 
 /**
- * Create and start an PXEService instance with the given AztecNode.
- * Returns a Promise that resolves to the started PXEService instance.
+ * Create and start an PXE instance with the given AztecNode.
+ * Returns a Promise that resolves to the started PXE instance.
  *
  * @param aztecNode - The AztecNode instance to be used by the server.
- * @param config - The PXE Service Config to use
+ * @param config - The PXE Config to use
  * @param
- * @returns A Promise that resolves to the started PXEService instance.
+ * @returns A Promise that resolves to the started PXE instance.
  */
-export async function createPXEService(
+export async function createPXE(
   aztecNode: AztecNode,
-  config: PXEServiceConfig,
+  config: PXEConfig,
   options: PXECreationOptions = { loggers: {} },
 ) {
+  const logSuffix =
+    typeof options.useLogSuffix === 'boolean'
+      ? options.useLogSuffix
+        ? randomBytes(3).toString('hex')
+        : undefined
+      : options.useLogSuffix;
+
   const l1Contracts = await aztecNode.getL1ContractAddresses();
   const configWithContracts = {
     ...config,
     l1Contracts,
-  } as PXEServiceConfig;
-
-  const store = await createStore(
-    'pxe_data',
-    configWithContracts,
-    options.loggers.store ?? createLogger('pxe:data:indexeddb'),
-  );
-
-  const simulationProvider = new WASMSimulator();
-  const prover =
-    options.prover ??
-    new BBWASMLazyPrivateKernelProver(simulationProvider, 16, options.loggers.prover ?? createLogger('bb:wasm:lazy'));
+  } as PXEConfig;
+
+  const loggers = options.loggers ?? {};
+
+  const storeLogger = loggers.store ? loggers.store : createLogger('pxe:data:idb' + (logSuffix ? `:${logSuffix}` : ''));
+
+  const store = options.store ?? (await createStore('pxe_data', configWithContracts, storeLogger));
+
+  const simulator = options.simulator ?? new WASMSimulator();
+  const proverLogger = loggers.prover
+    ? loggers.prover
+    : createLogger('pxe:bb:wasm:bundle' + (logSuffix ? `:${logSuffix}` : ''));
+
+  const prover = options.prover ?? new BBLazyPrivateKernelProver(simulator, proverLogger);
+
   const protocolContractsProvider = new LazyProtocolContractsProvider();
-  const pxe = await PXEService.create(
-    aztecNode,
-    store,
-    prover,
-    simulationProvider,
-    protocolContractsProvider,
-    config,
-    options.loggers.pxe ?? createLogger('pxe:service'),
-  );
+
+  const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));
+  const pxe = await PXE.create(aztecNode, store, prover, simulator, protocolContractsProvider, config, pxeLogger);
   return pxe;
 }

package/src/entrypoints/pxe_creation_options.ts

@@ -0,0 +1,12 @@
+import type { Logger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+import type { CircuitSimulator } from '@aztec/simulator/client';
+import type { PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
+
+export type PXECreationOptions = {
+  loggers?: { store?: Logger; pxe?: Logger; prover?: Logger };
+  useLogSuffix?: boolean | string;
+  prover?: PrivateKernelProver;
+  store?: AztecAsyncKVStore;
+  simulator?: CircuitSimulator;
+};

package/src/entrypoints/server/index.ts

@@ -1,6 +1,8 @@
-export * from '../../pxe_service/index.js';
-export * from '../../pxe_http/index.js';
+export * from '../../pxe.js';
 export * from '../../config/index.js';
+export * from '../../error_enriching.js';
 export * from '../../storage/index.js';
 export * from './utils.js';
-export { PXEOracleInterface } from '../../pxe_oracle_interface/index.js';
+export { NoteService } from '../../notes/note_service.js';
+export { ORACLE_VERSION } from '../../oracle_version.js';
+export { type PXECreationOptions } from '../pxe_creation_options.js';

package/src/entrypoints/server/utils.ts

@@ -1,71 +1,79 @@
-import { BBNativePrivateKernelProver } from '@aztec/bb-prover';
-import { BBWASMBundlePrivateKernelProver } from '@aztec/bb-prover/wasm/bundle';
-import { randomBytes } from '@aztec/foundation/crypto';
-import { createLogger } from '@aztec/foundation/log';
-import { createStore } from '@aztec/kv-store/lmdb-v2';
+import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { type Logger, createLogger } from '@aztec/foundation/log';
 import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import { type SimulationProvider, WASMSimulator } from '@aztec/simulator/client';
-import type { AztecNode, PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
+import {
+  type CircuitSimulator,
+  MemoryCircuitRecorder,
+  SimulatorRecorderWrapper,
+  WASMSimulator,
+} from '@aztec/simulator/client';
+import { FileCircuitRecorder } from '@aztec/simulator/testing';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
-import type { PXEServiceConfig } from '../../config/index.js';
-import { PXEService } from '../../pxe_service/pxe_service.js';
-import { PXE_DATA_SCHEMA_VERSION } from './index.js';
+import type { PXEConfig } from '../../config/index.js';
+import { PXE } from '../../pxe.js';
+import { PXE_DATA_SCHEMA_VERSION } from '../../storage/index.js';
+import type { PXECreationOptions } from '../pxe_creation_options.js';
 
-/**
- * Create and start an PXEService instance with the given AztecNode.
- * If no keyStore or database is provided, it will use KeyStore and MemoryDB as default values.
- * Returns a Promise that resolves to the started PXEService instance.
- *
- * @param aztecNode - The AztecNode instance to be used by the server.
- * @param config - The PXE Service Config to use
- * @param useLogSuffix - (Optional) Log suffix for PXE's logger.
- * @param proofCreator - An optional proof creator to use in place of any other configuration
- * @returns A Promise that resolves to the started PXEService instance.
- */
-export async function createPXEService(
+type PXEConfigWithoutDefaults = Omit<PXEConfig, 'l1Contracts' | 'l1ChainId' | 'l2BlockBatchSize' | 'rollupVersion'>;
+
+export async function createPXE(
   aztecNode: AztecNode,
-  config: PXEServiceConfig,
-  useLogSuffix: string | boolean | undefined = undefined,
-  proofCreator?: PrivateKernelProver,
+  config: PXEConfigWithoutDefaults,
+  options: PXECreationOptions = { loggers: {} },
 ) {
+  const recorder = process.env.CIRCUIT_RECORD_DIR
+    ? new FileCircuitRecorder(process.env.CIRCUIT_RECORD_DIR)
+    : new MemoryCircuitRecorder();
+  const simulator = new SimulatorRecorderWrapper(new WASMSimulator(), recorder);
+
   const logSuffix =
-    typeof useLogSuffix === 'boolean' ? (useLogSuffix ? randomBytes(3).toString('hex') : undefined) : useLogSuffix;
+    typeof options.useLogSuffix === 'boolean'
+      ? options.useLogSuffix
+        ? randomBytes(3).toString('hex')
+        : undefined
+      : options.useLogSuffix;
+  const loggers = options.loggers ?? {};
 
-  const l1Contracts = await aztecNode.getL1ContractAddresses();
-  const configWithContracts = {
+  const { l1ChainId, l1ContractAddresses: l1Contracts, rollupVersion } = await aztecNode.getNodeInfo();
+  const configWithContracts: PXEConfig = {
     ...config,
     l1Contracts,
-  } as PXEServiceConfig;
+    l1ChainId,
+    rollupVersion,
+    l2BlockBatchSize: 50,
+  };
 
-  const store = await createStore(
-    'pxe_data',
-    PXE_DATA_SCHEMA_VERSION,
-    configWithContracts,
-    createLogger('pxe:data:lmdb'),
-  );
+  if (!options.store) {
+    // TODO once https://github.com/AztecProtocol/aztec-packages/issues/13656 is fixed, we can remove this and always
+    // import the lmdb-v2 version
+    const { createStore } = await import('@aztec/kv-store/lmdb-v2');
+    const storeLogger = loggers.store
+      ? loggers.store
+      : createLogger('pxe:data:lmdb' + (logSuffix ? `:${logSuffix}` : ''));
+    options.store = await createStore('pxe_data', PXE_DATA_SCHEMA_VERSION, configWithContracts, storeLogger);
+  }
+  const proverLogger = loggers.prover
+    ? loggers.prover
+    : createLogger('pxe:bb:native' + (logSuffix ? `:${logSuffix}` : ''));
 
-  const simulationProvider = new WASMSimulator();
-  const prover = proofCreator ?? (await createProver(config, simulationProvider, logSuffix));
+  const prover = options.prover ?? createProver(simulator, proverLogger);
   const protocolContractsProvider = new BundledProtocolContractsProvider();
-  const pxe = await PXEService.create(
+
+  const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));
+  const pxe = await PXE.create(
     aztecNode,
-    store,
+    options.store,
     prover,
-    simulationProvider,
+    simulator,
     protocolContractsProvider,
-    config,
-    logSuffix,
+    configWithContracts,
+    pxeLogger,
   );
   return pxe;
 }
 
-function createProver(config: PXEServiceConfig, simulationProvider: SimulationProvider, logSuffix?: string) {
-  if (!config.bbBinaryPath || !config.bbWorkingDirectory) {
-    return new BBWASMBundlePrivateKernelProver(simulationProvider, 16);
-  } else {
-    const bbConfig = config as Required<Pick<PXEServiceConfig, 'bbBinaryPath' | 'bbWorkingDirectory'>> &
-      PXEServiceConfig;
-    const log = createLogger('pxe:bb-native-prover' + (logSuffix ? `:${logSuffix}` : ''));
-    return BBNativePrivateKernelProver.new({ bbSkipCleanup: false, ...bbConfig }, simulationProvider, log);
-  }
+function createProver(simulator: CircuitSimulator, logger?: Logger) {
+  return new BBBundlePrivateKernelProver(simulator, logger);
 }

package/src/{pxe_service/error_enriching.ts → error_enriching.ts}

@@ -1,12 +1,10 @@
-import { PUBLIC_DISPATCH_SELECTOR } from '@aztec/constants';
-import { Fr } from '@aztec/foundation/fields';
 import type { Logger } from '@aztec/foundation/log';
 import { resolveAssertionMessageFromRevertData, resolveOpcodeLocations } from '@aztec/simulator/client';
 import { FunctionSelector } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { type SimulationError, isNoirCallStackUnresolved } from '@aztec/stdlib/errors';
 
-import type { ContractDataProvider } from '../storage/contract_data_provider/contract_data_provider.js';
+import type { ContractDataProvider } from './storage/contract_data_provider/contract_data_provider.js';
 
 /**
  * Adds contract and function names to a simulation error, if they
@@ -18,35 +16,46 @@ export async function enrichSimulationError(
   contractDataProvider: ContractDataProvider,
   logger: Logger,
 ) {
-  // Maps contract addresses to the set of function names that were in error.
+  // Maps contract addresses to the set of function selectors that were in error.
   // Map and Set do reference equality for their keys instead of value equality, so we store the string
   // representation to get e.g. different contract address objects with the same address value to match.
-  const mentionedFunctions: Map<string, Set<string>> = new Map();
+  // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
+  const mentionedFunctions: Map<string, Set<FunctionSelector>> = new Map();
 
-  err.getCallStack().forEach(({ contractAddress, functionName }) => {
+  err.getCallStack().forEach(({ contractAddress, functionSelector }) => {
     if (!mentionedFunctions.has(contractAddress.toString())) {
       mentionedFunctions.set(contractAddress.toString(), new Set());
     }
-    mentionedFunctions.get(contractAddress.toString())!.add(functionName?.toString() ?? '');
+    if (functionSelector) {
+      mentionedFunctions.get(contractAddress.toString())!.add(functionSelector);
+    }
   });
 
   await Promise.all(
-    [...mentionedFunctions.entries()].map(async ([contractAddress, fnNames]) => {
+    [...mentionedFunctions.entries()].map(async ([contractAddress, fnSelectors]) => {
       const parsedContractAddress = AztecAddress.fromString(contractAddress);
       const contract = await contractDataProvider.getContract(parsedContractAddress);
       if (contract) {
         err.enrichWithContractName(parsedContractAddress, contract.name);
-        for (const fnName of fnNames) {
-          const functionArtifact = contract.functions.find(f => fnName === f.name);
-          if (functionArtifact) {
+        // Map from function selector to function name. It uses a stringified key for the same reason as mentionedFunctions.
+        const selectorToNameMap: Map<string, string> = new Map();
+        await Promise.all(
+          contract.functions.map(async fn => {
+            const selector = await FunctionSelector.fromNameAndParameters(fn);
+            selectorToNameMap.set(selector.toString(), fn.name);
+          }),
+        );
+
+        for (const fnSelector of fnSelectors) {
+          if (selectorToNameMap.has(fnSelector.toString())) {
             err.enrichWithFunctionName(
               parsedContractAddress,
-              await FunctionSelector.fromNameAndParameters(functionArtifact),
-              functionArtifact.name,
+              fnSelector,
+              selectorToNameMap.get(fnSelector.toString())!,
             );
           } else {
             logger.warn(
-              `Could not find function artifact in contract ${contract.name} for function '${fnName}' when enriching error callstack`,
+              `Could not find function artifact in contract ${contract.name} for function '${fnSelector}' when enriching error callstack`,
             );
           }
         }
@@ -66,24 +75,24 @@ export async function enrichPublicSimulationError(
 ) {
   const callStack = err.getCallStack();
   const originalFailingFunction = callStack[callStack.length - 1];
-  // TODO(https://github.com/AztecProtocol/aztec-packages/issues/8985): Properly fix this.
-  // 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 contractDataProvider.getFunctionArtifact(
-    originalFailingFunction.contractAddress,
-    FunctionSelector.fromField(new Fr(PUBLIC_DISPATCH_SELECTOR)),
-  );
+  if (!originalFailingFunction) {
+    throw new Error(`Original failing function not found when enriching public simulation, missing callstack`);
+  }
+
+  const artifact = await contractDataProvider.getPublicFunctionArtifact(originalFailingFunction.contractAddress);
+  if (!artifact) {
+    throw new Error(
+      `Artifact not found when enriching public simulation error. Contract address: ${originalFailingFunction.contractAddress}.`,
+    );
+  }
+
   const assertionMessage = resolveAssertionMessageFromRevertData(err.revertData, artifact);
   if (assertionMessage) {
     err.setOriginalMessage(err.getOriginalMessage() + `${assertionMessage}`);
   }
 
-  const debugInfo = await contractDataProvider.getFunctionDebugMetadata(
-    originalFailingFunction.contractAddress,
-    FunctionSelector.fromField(new Fr(PUBLIC_DISPATCH_SELECTOR)),
-  );
+  const debugInfo = await contractDataProvider.getPublicFunctionDebugMetadata(originalFailingFunction.contractAddress);
 
   const noirCallStack = err.getNoirCallStack();
   if (debugInfo) {
@@ -91,7 +100,7 @@ export async function enrichPublicSimulationError(
       try {
         // Public functions are simulated as a single Brillig entry point.
         // Thus, we can safely assume here that the Brillig function id is `0`.
-        const parsedCallStack = resolveOpcodeLocations(noirCallStack, debugInfo, 0);
+        const parsedCallStack = resolveOpcodeLocations(noirCallStack, debugInfo.debugSymbols, debugInfo.files, 0);
         err.setNoirCallStack(parsedCallStack);
       } catch (err) {
         logger.warn(