@aztec/simulator 0.0.0-test.0 → 0.0.1-commit.001888fc

package/src/public/fixtures/public_tx_simulation_tester.ts

@@ -1,53 +1,112 @@
-import { PUBLIC_DISPATCH_SELECTOR } from '@aztec/constants';
-import { Fr } from '@aztec/foundation/fields';
-import { type ContractArtifact, FunctionSelector, encodeArguments } from '@aztec/stdlib/abi';
+import {
+  DEFAULT_TEARDOWN_DA_GAS_LIMIT,
+  DEFAULT_TEARDOWN_L2_GAS_LIMIT,
+  PUBLIC_TX_L2_GAS_OVERHEAD,
+  TX_DA_GAS_OVERHEAD,
+} from '@aztec/constants';
+import { asyncMap } from '@aztec/foundation/async-map';
+import { BlockNumber } from '@aztec/foundation/branded-types';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { type ContractArtifact, encodeArguments } from '@aztec/stdlib/abi';
+import { PublicSimulatorConfig, type PublicTxResult } from '@aztec/stdlib/avm';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { GasFees } from '@aztec/stdlib/gas';
+import { Gas, GasFees } from '@aztec/stdlib/gas';
 import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
-import { PublicExecutionRequest, type Tx } from '@aztec/stdlib/tx';
-import { CallContext, GlobalVariables } from '@aztec/stdlib/tx';
+import { PublicCallRequest } from '@aztec/stdlib/kernel';
+import { GlobalVariables, PublicCallRequestWithCalldata, type Tx } from '@aztec/stdlib/tx';
 import { NativeWorldStateService } from '@aztec/world-state';
 
 import { BaseAvmSimulationTester } from '../avm/fixtures/base_avm_simulation_tester.js';
-import { getContractFunctionArtifact, getFunctionSelector } from '../avm/fixtures/index.js';
-import { SimpleContractDataSource } from '../avm/fixtures/simple_contract_data_source.js';
-import { WorldStateDB } from '../public_db_sources.js';
-import { type PublicTxResult, PublicTxSimulator } from '../public_tx_simulator/public_tx_simulator.js';
-import { createTxForPublicCalls } from './index.js';
+import {
+  DEFAULT_BLOCK_NUMBER,
+  DEFAULT_TIMESTAMP,
+  getContractFunctionAbi,
+  getFunctionSelector,
+} from '../avm/fixtures/utils.js';
+import { PublicContractsDB } from '../public_db_sources.js';
+import { MeasuredCppPublicTxSimulator } from '../public_tx_simulator/cpp_public_tx_simulator.js';
+import { MeasuredCppVsTsPublicTxSimulator } from '../public_tx_simulator/cpp_vs_ts_public_tx_simulator.js';
+import type { MeasuredPublicTxSimulatorInterface } from '../public_tx_simulator/public_tx_simulator_interface.js';
+import { TestExecutorMetrics } from '../test_executor_metrics.js';
+import { SimpleContractDataSource } from './simple_contract_data_source.js';
+import { type TestPrivateInsertions, createTxForPublicCalls } from './utils.js';
 
-const TIMESTAMP = new Fr(99833);
 const DEFAULT_GAS_FEES = new GasFees(2, 3);
-export const DEFAULT_BLOCK_NUMBER = 42;
 
 export type TestEnqueuedCall = {
+  sender?: AztecAddress;
   address: AztecAddress;
-  fnName: string;
+  fnName?: string;
   args: any[];
   isStaticCall?: boolean;
   contractArtifact?: ContractArtifact;
 };
 
+const defaultConfig: PublicSimulatorConfig = PublicSimulatorConfig.from({
+  skipFeeEnforcement: false,
+  collectCallMetadata: true,
+  collectDebugLogs: true,
+  collectHints: false,
+  collectPublicInputs: false,
+  collectStatistics: false,
+});
+
+/**
+ * Factory type for creating a MeasuredPublicTxSimulatorInterface.
+ */
+export type MeasuredSimulatorFactory = (
+  merkleTree: MerkleTreeWriteOperations,
+  contractsDB: PublicContractsDB,
+  globals: GlobalVariables,
+  metrics: TestExecutorMetrics,
+  config: PublicSimulatorConfig,
+) => MeasuredPublicTxSimulatorInterface;
+
 /**
  * A test class that extends the BaseAvmSimulationTester to enable real-app testing of the PublicTxSimulator.
  * It provides an interface for simulating one transaction at a time and maintains state between subsequent
  * transactions.
  */
 export class PublicTxSimulationTester extends BaseAvmSimulationTester {
-  private txCount = 0;
+  protected txCount: number = 0;
+  private simulator: MeasuredPublicTxSimulatorInterface;
+  private metricsPrefix?: string;
 
   constructor(
-    private worldStateDB: WorldStateDB,
+    merkleTree: MerkleTreeWriteOperations,
     contractDataSource: SimpleContractDataSource,
-    merkleTrees: MerkleTreeWriteOperations,
+    globals: GlobalVariables = defaultGlobals(),
+    private metrics: TestExecutorMetrics = new TestExecutorMetrics(),
+    simulatorFactory?: MeasuredSimulatorFactory,
+    config: PublicSimulatorConfig = defaultConfig,
   ) {
-    super(contractDataSource, merkleTrees);
+    super(contractDataSource, merkleTree);
+
+    const contractsDB = new PublicContractsDB(contractDataSource);
+    if (simulatorFactory) {
+      this.simulator = simulatorFactory(merkleTree, contractsDB, globals, this.metrics, config);
+    } else {
+      this.simulator = new MeasuredCppPublicTxSimulator(merkleTree, contractsDB, globals, this.metrics, config);
+    }
   }
 
-  public static async create(): Promise<PublicTxSimulationTester> {
+  public static async create(
+    worldStateService: NativeWorldStateService, // make sure to close this later
+    globals: GlobalVariables = defaultGlobals(),
+    metrics: TestExecutorMetrics = new TestExecutorMetrics(),
+    useCppSimulator = false,
+    config: PublicSimulatorConfig = defaultConfig,
+  ): Promise<PublicTxSimulationTester> {
     const contractDataSource = new SimpleContractDataSource();
-    const merkleTrees = await (await NativeWorldStateService.tmp()).fork();
-    const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
-    return new PublicTxSimulationTester(worldStateDB, contractDataSource, merkleTrees);
+    const merkleTree = await worldStateService.fork();
+    const simulatorFactory: MeasuredSimulatorFactory = useCppSimulator
+      ? (mt, cdb, g, m, c) => new MeasuredCppPublicTxSimulator(mt, cdb, g, m, c)
+      : (mt, cdb, g, m, c) => new MeasuredCppVsTsPublicTxSimulator(mt, cdb, g, m, c);
+    return new PublicTxSimulationTester(merkleTree, contractDataSource, globals, metrics, simulatorFactory, config);
+  }
+
+  public setMetricsPrefix(prefix: string) {
+    this.metricsPrefix = prefix;
   }
 
   public async createTx(
@@ -57,69 +116,34 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     teardownCall?: TestEnqueuedCall,
     feePayer: AztecAddress = sender,
     /* need some unique first nullifier for note-nonce computations */
-    firstNullifier = new Fr(420000 + this.txCount++),
+    privateInsertions: TestPrivateInsertions = { nonRevertible: { nullifiers: [new Fr(420000 + this.txCount)] } },
+    gasLimits?: Gas,
   ): Promise<Tx> {
-    const setupExecutionRequests: PublicExecutionRequest[] = [];
-    for (let i = 0; i < setupCalls.length; i++) {
-      const address = setupCalls[i].address;
-      const contractArtifact =
-        setupCalls[i].contractArtifact || (await this.contractDataSource.getContractArtifact(address));
-      if (!contractArtifact) {
-        throw new Error(`Contract artifact not found for address: ${address}`);
-      }
-      const req = await executionRequestForCall(
-        contractArtifact,
-        sender,
-        address,
-        setupCalls[i].fnName,
-        setupCalls[i].args,
-        setupCalls[i].isStaticCall,
-      );
-      setupExecutionRequests.push(req);
-    }
-    const appExecutionRequests: PublicExecutionRequest[] = [];
-    for (let i = 0; i < appCalls.length; i++) {
-      const address = appCalls[i].address;
-      const contractArtifact =
-        appCalls[i].contractArtifact || (await this.contractDataSource.getContractArtifact(address));
-      if (!contractArtifact) {
-        throw new Error(`Contract artifact not found for address: ${address}`);
-      }
-      const req = await executionRequestForCall(
-        contractArtifact,
-        sender,
-        address,
-        appCalls[i].fnName,
-        appCalls[i].args,
-        appCalls[i].isStaticCall,
-      );
-      appExecutionRequests.push(req);
-    }
-
-    let teardownExecutionRequest: PublicExecutionRequest | undefined = undefined;
-    if (teardownCall) {
-      const address = teardownCall.address;
-      const contractArtifact =
-        teardownCall.contractArtifact || (await this.contractDataSource.getContractArtifact(address));
-      if (!contractArtifact) {
-        throw new Error(`Contract artifact not found for address: ${address}`);
-      }
-      teardownExecutionRequest = await executionRequestForCall(
-        contractArtifact,
-        sender,
-        address,
-        teardownCall.fnName,
-        teardownCall.args,
-        teardownCall.isStaticCall,
-      );
-    }
+    const setupCallRequests = await asyncMap(setupCalls, call =>
+      this.#createPubicCallRequestForCall(call, call.sender ?? sender),
+    );
+    const appCallRequests = await asyncMap(appCalls, call =>
+      this.#createPubicCallRequestForCall(call, call.sender ?? sender),
+    );
+    const teardownCallRequest = teardownCall
+      ? await this.#createPubicCallRequestForCall(teardownCall, teardownCall.sender ?? sender)
+      : undefined;
 
-    return await createTxForPublicCalls(
-      firstNullifier,
-      setupExecutionRequests,
-      appExecutionRequests,
-      teardownExecutionRequest,
+    this.txCount++;
+    return createTxForPublicCalls(
+      privateInsertions,
+      setupCallRequests,
+      appCallRequests,
+      teardownCallRequest,
       feePayer,
+      /*gasUsedByPrivate*/ teardownCall
+        ? new Gas(
+            DEFAULT_TEARDOWN_DA_GAS_LIMIT + TX_DA_GAS_OVERHEAD,
+            DEFAULT_TEARDOWN_L2_GAS_LIMIT + PUBLIC_TX_L2_GAS_OVERHEAD,
+          )
+        : new Gas(TX_DA_GAS_OVERHEAD, PUBLIC_TX_L2_GAS_OVERHEAD),
+      defaultGlobals(),
+      gasLimits,
     );
   }
 
@@ -130,50 +154,181 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     teardownCall?: TestEnqueuedCall,
     feePayer: AztecAddress = sender,
     /* need some unique first nullifier for note-nonce computations */
-    firstNullifier = new Fr(420000 + this.txCount++),
-    globals = defaultGlobals(),
+    privateInsertions?: TestPrivateInsertions,
+    txLabel: string = 'unlabeledTx',
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
-    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, firstNullifier);
+    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions, gasLimits);
 
     await this.setFeePayerBalance(feePayer);
 
-    const simulator = new PublicTxSimulator(this.merkleTrees, this.worldStateDB, globals, /*doMerkleOperations=*/ true);
+    const txLabelWithCount = `${txLabel}/${this.txCount - 1}`;
+    const fullTxLabel = this.metricsPrefix ? `${this.metricsPrefix}/${txLabelWithCount}` : txLabelWithCount;
+
+    if (!this.simulator) {
+      throw new Error(
+        'No simulator configured. Pass a simulatorFactory to the constructor or use PublicTxSimulationTester.create()',
+      );
+    }
+    const avmResult = await this.simulator.simulate(tx, fullTxLabel);
 
-    const startTime = performance.now();
-    const avmResult = await simulator.simulate(tx);
-    const endTime = performance.now();
-    this.logger.debug(`Public transaction simulation took ${endTime - startTime}ms`);
+    await this.#recordBytecodeSizes(fullTxLabel, [...setupCalls, ...appCalls, ...(teardownCall ? [teardownCall] : [])]);
+
+    // Something like this is often useful for debugging:
+    //if (avmResult.revertReason) {
+    //  // resolve / enrich revert reason
+    //  const lastAppCall = appCalls[appCalls.length - 1];
+
+    //  const contractArtifact =
+    //    lastAppCall.contractArtifact || (await this.contractDataSource.getContractArtifact(lastAppCall.address));
+    //  const fnAbi = getContractFunctionAbi(lastAppCall.fnName, contractArtifact!);
+    //  const revertReason = resolveAssertionMessageFromRevertData(avmResult.revertReason.revertData, fnAbi!);
+    //  this.logger.debug(`Revert reason: ${revertReason}`);
+    //}
 
     return avmResult;
   }
-}
 
-async function executionRequestForCall(
-  contractArtifact: ContractArtifact,
-  sender: AztecAddress,
-  address: AztecAddress,
-  fnName: string,
-  args: Fr[] = [],
-  isStaticCall: boolean = false,
-): Promise<PublicExecutionRequest> {
-  const fnSelector = await getFunctionSelector(fnName, contractArtifact);
-  const fnAbi = getContractFunctionArtifact(fnName, contractArtifact);
-  const encodedArgs = encodeArguments(fnAbi!, args);
-  const calldata = [fnSelector.toField(), ...encodedArgs];
-
-  const callContext = new CallContext(
-    sender,
-    address,
-    /*selector=*/ new FunctionSelector(PUBLIC_DISPATCH_SELECTOR),
-    isStaticCall,
-  );
-  return new PublicExecutionRequest(callContext, calldata);
+  /**
+   * Just simulate the transaction and return the result.
+   *
+   * This wrapper around simulation allows for easy labeling of a TX
+   * which is especially useful when reporting benchmarks or metrics.
+   */
+  public async simulateTxWithLabel(
+    txLabel: string,
+    sender: AztecAddress,
+    setupCalls?: TestEnqueuedCall[],
+    appCalls?: TestEnqueuedCall[],
+    teardownCall?: TestEnqueuedCall,
+    feePayer?: AztecAddress,
+    privateInsertions?: TestPrivateInsertions,
+    gasLimits?: Gas,
+  ): Promise<PublicTxResult> {
+    return await this.simulateTx(
+      sender,
+      setupCalls,
+      appCalls,
+      teardownCall,
+      feePayer,
+      privateInsertions,
+      txLabel,
+      gasLimits,
+    );
+  }
+
+  /**
+   * Execute a transaction and return the result.
+   *
+   * This function can be (it is) overridden by a subclass (AvmProvingTester)
+   * to do more work (like prove and verify) while still reusing existing
+   * test fixtures (like amm_test). That is why it is not named "simulate*".
+   */
+  public async executeTxWithLabel(
+    txLabel: string,
+    sender: AztecAddress,
+    setupCalls?: TestEnqueuedCall[],
+    appCalls?: TestEnqueuedCall[],
+    teardownCall?: TestEnqueuedCall,
+    feePayer?: AztecAddress,
+    privateInsertions?: TestPrivateInsertions,
+    gasLimits?: Gas,
+  ): Promise<PublicTxResult> {
+    return await this.simulateTxWithLabel(
+      txLabel,
+      sender,
+      setupCalls,
+      appCalls,
+      teardownCall,
+      feePayer,
+      privateInsertions,
+      gasLimits,
+    );
+  }
+
+  public prettyPrintMetrics() {
+    this.metrics.prettyPrint();
+  }
+
+  /**
+   * Cancel the current simulation if one is in progress.
+   * This signals the underlying simulator (e.g., C++) to stop at the next safe point.
+   * Safe to call even if no simulation is in progress.
+   *
+   * @param waitTimeoutMs - If provided, wait up to this many ms for the simulation to actually stop.
+   */
+  public async cancel(waitTimeoutMs?: number): Promise<void> {
+    await this.simulator.cancel?.(waitTimeoutMs);
+  }
+
+  /**
+   * Get the underlying simulator for advanced test scenarios.
+   * Use this when you need direct control over simulation (e.g., for testing cancellation).
+   */
+  public getSimulator(): MeasuredPublicTxSimulatorInterface {
+    return this.simulator;
+  }
+
+  async #createPubicCallRequestForCall(
+    call: TestEnqueuedCall,
+    sender: AztecAddress,
+  ): Promise<PublicCallRequestWithCalldata> {
+    const address = call.address;
+    const contractArtifact = call.contractArtifact || (await this.contractDataSource.getContractArtifact(address));
+    if (!contractArtifact) {
+      throw new Error(`Contract artifact not found for address: ${address}`);
+    }
+
+    let calldata: Fr[] = [];
+    if (!call.fnName) {
+      this.logger.debug(
+        `No function name specified for call to contract ${call.address.toString()}. Assuming this is a custom bytecode with no public_dispatch function.`,
+      );
+      this.logger.debug(`Not using ABI to encode arguments. Not prepending fn selector to calldata.`);
+      try {
+        calldata = call.args.map(arg => new Fr(arg));
+      } catch (error) {
+        this.logger.warn(`Tried assuming that all arguments are Field-like. Failed. Error: ${error}`);
+        throw error;
+      }
+    } else {
+      const fnSelector = await getFunctionSelector(call.fnName, contractArtifact);
+      const fnAbi = getContractFunctionAbi(call.fnName, contractArtifact)!;
+      const encodedArgs = encodeArguments(fnAbi, call.args);
+      calldata = [fnSelector.toField(), ...encodedArgs];
+    }
+    const isStaticCall = call.isStaticCall ?? false;
+    const request = await PublicCallRequest.fromCalldata(sender, address, isStaticCall, calldata);
+
+    return new PublicCallRequestWithCalldata(request, calldata);
+  }
+
+  // WARNING: Deduplicates by artifact name, so two different artifacts with the same name
+  // in a single tx would only record the first one's bytecode size.
+  async #recordBytecodeSizes(txLabel: string, calls: TestEnqueuedCall[]) {
+    const seenArtifactNames = new Set<string>();
+    for (const call of calls) {
+      const artifact = await this.contractDataSource.getContractArtifact(call.address);
+      if (!artifact || seenArtifactNames.has(artifact.name)) {
+        continue;
+      }
+      seenArtifactNames.add(artifact.name);
+      const instance = await this.contractDataSource.getContract(call.address);
+      if (!instance) {
+        continue;
+      }
+      const contractClass = await this.contractDataSource.getContractClass(instance.currentContractClassId);
+      if (contractClass) {
+        this.metrics.recordBytecodeSize(txLabel, artifact.name, contractClass.packedBytecode.length);
+      }
+    }
+  }
 }
 
-function defaultGlobals() {
+export function defaultGlobals() {
   const globals = GlobalVariables.empty();
-  globals.timestamp = TIMESTAMP;
+  globals.timestamp = DEFAULT_TIMESTAMP;
   globals.gasFees = DEFAULT_GAS_FEES; // apply some nonzero default gas fees
-  globals.blockNumber = new Fr(DEFAULT_BLOCK_NUMBER);
+  globals.blockNumber = BlockNumber(DEFAULT_BLOCK_NUMBER);
   return globals;
 }

package/src/public/{avm/fixtures → fixtures}/simple_contract_data_source.ts

@@ -1,15 +1,11 @@
-import type { Fr } from '@aztec/foundation/fields';
+import type { BlockNumber } from '@aztec/foundation/branded-types';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { createLogger } from '@aztec/foundation/log';
-import { type ContractArtifact, FunctionSelector } from '@aztec/stdlib/abi';
+import type { ContractArtifact, FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
-import type {
-  ContractClassPublic,
-  ContractDataSource,
-  ContractInstanceWithAddress,
-  PublicFunction,
-} from '@aztec/stdlib/contract';
+import type { ContractClassPublic, ContractDataSource, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 
-import { PUBLIC_DISPATCH_FN_NAME } from './index.js';
+import { getFunctionSelector } from '../avm/fixtures/utils.js';
 
 /**
  * This class is used during public/avm testing to function as a database of
@@ -27,6 +23,8 @@ export class SimpleContractDataSource implements ContractDataSource {
   private contractInstances: Map<string, ContractInstanceWithAddress> = new Map();
   // maps contract instance address to address
   private contractArtifacts: Map<string, ContractArtifact> = new Map();
+  // maps `${classID}:${fnSelector}` to name
+  private debugFunctionName: Map<string, string> = new Map();
 
   /////////////////////////////////////////////////////////////
   // Helper functions not in the contract data source interface
@@ -39,22 +37,30 @@ export class SimpleContractDataSource implements ContractDataSource {
     contractClass: ContractClassPublic,
     contractInstance: ContractInstanceWithAddress,
   ) {
-    this.addContractArtifact(contractClass.id, contractArtifact);
+    await this.addContractArtifact(contractClass.id, contractArtifact);
     await this.addContractClass(contractClass);
     await this.addContractInstance(contractInstance);
   }
 
-  addContractArtifact(classId: Fr, artifact: ContractArtifact): void {
+  async addContractArtifact(classId: Fr, artifact: ContractArtifact) {
     this.contractArtifacts.set(classId.toString(), artifact);
+    const classIdStr = classId.toString();
+    const publicFns = artifact.nonDispatchPublicFunctions;
+    if (publicFns.length !== 0) {
+      for (const fn of publicFns) {
+        const actualFnName = `${fn.name}`;
+        const fnSelector = await getFunctionSelector(actualFnName, artifact);
+        const key = `${classIdStr}:${fnSelector.toString()}`;
+
+        const longFnName = `${artifact.name}.${actualFnName}`;
+        this.debugFunctionName.set(key, longFnName);
+      }
+    }
   }
 
   /////////////////////////////////////////////////////////////
   // ContractDataSource function implementations
-  getPublicFunction(_address: AztecAddress, _selector: FunctionSelector): Promise<PublicFunction> {
-    throw new Error('Method not implemented.');
-  }
-
-  getBlockNumber(): Promise<number> {
+  getBlockNumber(): Promise<BlockNumber> {
     throw new Error('Method not implemented.');
   }
 
@@ -82,14 +88,25 @@ export class SimpleContractDataSource implements ContractDataSource {
     }
     this.logger.debug(`Retrieved contract artifact for address: ${address}`);
     this.logger.debug(`Contract class ID: ${contractInstance.currentContractClassId}`);
-    return Promise.resolve(this.contractArtifacts.get(contractInstance!.currentContractClassId.toString()));
+    return this.contractArtifacts.get(contractInstance!.currentContractClassId.toString());
   }
 
-  getContractFunctionName(_address: AztecAddress, _selector: FunctionSelector): Promise<string> {
-    return Promise.resolve(PUBLIC_DISPATCH_FN_NAME);
+  async getDebugFunctionName(address: AztecAddress, selector: FunctionSelector): Promise<string | undefined> {
+    const contractInstance = await this.getContract(address);
+    if (!contractInstance) {
+      this.logger.warn(`Couldn't get fn name for debugging. Contract not in tester's ContractDataSource.`);
+      return undefined;
+    }
+    const key = `${contractInstance.currentContractClassId.toString()}:${selector.toString()}`;
+    const fnName = this.debugFunctionName.get(key);
+    if (!fnName) {
+      this.logger.warn(`Couldn't get fn name for debugging...`);
+      return undefined;
+    }
+    return fnName;
   }
 
-  registerContractFunctionSignatures(_address: AztecAddress, _signatures: string[]): Promise<void> {
+  registerContractFunctionSignatures(_signatures: string[]): Promise<void> {
     return Promise.resolve();
   }