@aztec/simulator 0.0.1-commit.fce3e4f → 0.0.1-commit.ffe5b04ea

package/src/public/fixtures/public_tx_simulation_tester.ts

@@ -1,6 +1,12 @@
-import { DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT } from '@aztec/constants';
+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 { Fr } from '@aztec/foundation/fields';
+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';
@@ -19,7 +25,7 @@ import {
 } from '../avm/fixtures/utils.js';
 import { PublicContractsDB } from '../public_db_sources.js';
 import { MeasuredCppPublicTxSimulator } from '../public_tx_simulator/cpp_public_tx_simulator.js';
-import { MeasuredPublicTxSimulator } from '../public_tx_simulator/measured_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';
@@ -30,7 +36,7 @@ const DEFAULT_GAS_FEES = new GasFees(2, 3);
 export type TestEnqueuedCall = {
   sender?: AztecAddress;
   address: AztecAddress;
-  fnName: string;
+  fnName?: string;
   args: any[];
   isStaticCall?: boolean;
   contractArtifact?: ContractArtifact;
@@ -41,9 +47,21 @@ const defaultConfig: PublicSimulatorConfig = PublicSimulatorConfig.from({
   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
@@ -59,15 +77,17 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     contractDataSource: SimpleContractDataSource,
     globals: GlobalVariables = defaultGlobals(),
     private metrics: TestExecutorMetrics = new TestExecutorMetrics(),
-    useCppSimulator: boolean = false,
+    simulatorFactory?: MeasuredSimulatorFactory,
     config: PublicSimulatorConfig = defaultConfig,
   ) {
     super(contractDataSource, merkleTree);
 
     const contractsDB = new PublicContractsDB(contractDataSource);
-    this.simulator = useCppSimulator
-      ? new MeasuredCppPublicTxSimulator(merkleTree, contractsDB, globals, this.metrics, config)
-      : new MeasuredPublicTxSimulator(merkleTree, contractsDB, globals, this.metrics, config);
+    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(
@@ -79,7 +99,10 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
   ): Promise<PublicTxSimulationTester> {
     const contractDataSource = new SimpleContractDataSource();
     const merkleTree = await worldStateService.fork();
-    return new PublicTxSimulationTester(merkleTree, contractDataSource, globals, metrics, useCppSimulator, config);
+    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) {
@@ -94,6 +117,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     feePayer: AztecAddress = sender,
     /* need some unique first nullifier for note-nonce computations */
     privateInsertions: TestPrivateInsertions = { nonRevertible: { nullifiers: [new Fr(420000 + this.txCount)] } },
+    gasLimits?: Gas,
   ): Promise<Tx> {
     const setupCallRequests = await asyncMap(setupCalls, call =>
       this.#createPubicCallRequestForCall(call, call.sender ?? sender),
@@ -113,9 +137,13 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       teardownCallRequest,
       feePayer,
       /*gasUsedByPrivate*/ teardownCall
-        ? new Gas(DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT)
-        : Gas.empty(),
+        ? 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,
     );
   }
 
@@ -128,16 +156,24 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     /* need some unique first nullifier for note-nonce computations */
     privateInsertions?: TestPrivateInsertions,
     txLabel: string = 'unlabeledTx',
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
-    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions);
+    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions, gasLimits);
 
     await this.setFeePayerBalance(feePayer);
 
     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);
 
+    await this.#recordBytecodeSizes(fullTxLabel, [...setupCalls, ...appCalls, ...(teardownCall ? [teardownCall] : [])]);
+
     // Something like this is often useful for debugging:
     //if (avmResult.revertReason) {
     //  // resolve / enrich revert reason
@@ -167,8 +203,18 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     teardownCall?: TestEnqueuedCall,
     feePayer?: AztecAddress,
     privateInsertions?: TestPrivateInsertions,
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
-    return await this.simulateTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions, txLabel);
+    return await this.simulateTx(
+      sender,
+      setupCalls,
+      appCalls,
+      teardownCall,
+      feePayer,
+      privateInsertions,
+      txLabel,
+      gasLimits,
+    );
   }
 
   /**
@@ -186,6 +232,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     teardownCall?: TestEnqueuedCall,
     feePayer?: AztecAddress,
     privateInsertions?: TestPrivateInsertions,
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
     return await this.simulateTxWithLabel(
       txLabel,
@@ -195,6 +242,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       teardownCall,
       feePayer,
       privateInsertions,
+      gasLimits,
     );
   }
 
@@ -202,6 +250,25 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     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,
@@ -212,21 +279,56 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       throw new Error(`Contract artifact not found for address: ${address}`);
     }
 
-    const fnSelector = await getFunctionSelector(call.fnName, contractArtifact);
-    const fnAbi = getContractFunctionAbi(call.fnName, contractArtifact)!;
-    const encodedArgs = encodeArguments(fnAbi, call.args);
-    const calldata = [fnSelector.toField(), ...encodedArgs];
+    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);
+      }
+    }
+  }
 }
 
 export function defaultGlobals() {
   const globals = GlobalVariables.empty();
   globals.timestamp = DEFAULT_TIMESTAMP;
   globals.gasFees = DEFAULT_GAS_FEES; // apply some nonzero default gas fees
-  globals.blockNumber = DEFAULT_BLOCK_NUMBER;
+  globals.blockNumber = BlockNumber(DEFAULT_BLOCK_NUMBER);
   return globals;
 }

package/src/public/fixtures/simple_contract_data_source.ts

@@ -1,4 +1,5 @@
-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 { AztecAddress } from '@aztec/stdlib/aztec-address';
@@ -59,7 +60,7 @@ export class SimpleContractDataSource implements ContractDataSource {
 
   /////////////////////////////////////////////////////////////
   // ContractDataSource function implementations
-  getBlockNumber(): Promise<number> {
+  getBlockNumber(): Promise<BlockNumber> {
     throw new Error('Method not implemented.');
   }
 
@@ -90,19 +91,17 @@ export class SimpleContractDataSource implements ContractDataSource {
     return this.contractArtifacts.get(contractInstance!.currentContractClassId.toString());
   }
 
-  async getDebugFunctionName(address: AztecAddress, selector: FunctionSelector): Promise<string> {
+  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. Using selector:${selector} instead...`,
-      );
-      return `selector:${selector.toString()}`;
+      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. Using selector:${selector} instead...`);
-      return selector.toString();
+      this.logger.warn(`Couldn't get fn name for debugging...`);
+      return undefined;
     }
     return fnName;
   }

package/src/public/fixtures/token_test.ts

@@ -1,4 +1,4 @@
-import { Fr } from '@aztec/foundation/fields';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import type { Logger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import type { ContractArtifact } from '@aztec/stdlib/abi';
@@ -7,11 +7,17 @@ import type { ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 
 import { PublicTxSimulationTester } from './public_tx_simulation_tester.js';
 
+export type TokenTestOptions = {
+  /** Skip return value assertions in balance checks (useful for benchmarking with collectCallMetadata=false) */
+  skipReturnValueAssertions?: boolean;
+};
+
 export async function tokenTest(
   tester: PublicTxSimulationTester,
   logger: Logger,
   tokenArtifact: ContractArtifact,
   expectToBeTrue: (x: boolean) => void,
+  options: TokenTestOptions = {},
 ) {
   const timer = new Timer();
 
@@ -36,7 +42,7 @@ export async function tokenTest(
     ],
   );
   expectToBeTrue(mintResult.revertCode.isOK());
-  await checkBalance(tester, token, sender, sender, mintAmount, expectToBeTrue);
+  await checkBalance(tester, token, sender, sender, mintAmount, expectToBeTrue, options);
 
   const authwitNonce = new Fr(0);
   const transferAmount = 50n;
@@ -54,8 +60,8 @@ export async function tokenTest(
     ],
   );
   expectToBeTrue(transferResult.revertCode.isOK());
-  await checkBalance(tester, token, sender, receiver, mintAmount - transferAmount, expectToBeTrue);
-  await checkBalance(tester, token, sender, receiver, transferAmount, expectToBeTrue);
+  await checkBalance(tester, token, sender, sender, mintAmount - transferAmount, expectToBeTrue, options);
+  await checkBalance(tester, token, sender, receiver, transferAmount, expectToBeTrue, options);
 
   // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
   const burnResult = await tester.executeTxWithLabel(
@@ -71,7 +77,7 @@ export async function tokenTest(
     ],
   );
   expectToBeTrue(burnResult.revertCode.isOK());
-  await checkBalance(tester, token, sender, receiver, 0n, expectToBeTrue);
+  await checkBalance(tester, token, sender, receiver, 0n, expectToBeTrue, options);
 
   logger.info(`TokenContract test took ${timer.ms()}ms\n`);
 }
@@ -116,6 +122,7 @@ async function checkBalance(
   account: AztecAddress,
   expectedBalance: bigint,
   expectToBeTrue: (x: boolean) => void,
+  options: TokenTestOptions = {},
 ) {
   // Strictly simulate this! No need to "execute" (aka prove if using AvmProvingTester subclass).
   const balResult = await tester.simulateTxWithLabel(
@@ -132,8 +139,10 @@ async function checkBalance(
     ],
   );
   expectToBeTrue(balResult.revertCode.isOK());
-  // should be 1 call with 1 return value that is expectedBalance
-  expectToBeTrue(balResult.appLogicReturnValues.length == 1);
-  expectToBeTrue(balResult.appLogicReturnValues[0].values!.length == 1);
-  expectToBeTrue(balResult.appLogicReturnValues[0].values![0].toBigInt() == expectedBalance);
+  if (!options.skipReturnValueAssertions) {
+    // should be 1 call with 1 return value that is expectedBalance
+    const appLogicReturnValues = balResult.getAppLogicReturnValues();
+    expectToBeTrue(appLogicReturnValues.length === 1);
+    expectToBeTrue(appLogicReturnValues[0].values?.[0]?.toBigInt() === expectedBalance);
+  }
 }

package/src/public/fixtures/utils.ts

@@ -9,7 +9,7 @@ import {
   PRIVATE_LOG_SIZE_IN_FIELDS,
 } from '@aztec/constants';
 import { padArrayEnd } from '@aztec/foundation/collection';
-import { Fr } from '@aztec/foundation/fields';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { CONTRACT_INSTANCE_PUBLISHED_EVENT_TAG } from '@aztec/protocol-contracts';
 import { bufferAsFields } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
@@ -62,13 +62,14 @@ export async function createTxForPublicCalls(
   feePayer = AztecAddress.zero(),
   gasUsedByPrivate: Gas = Gas.empty(),
   globals: GlobalVariables = GlobalVariables.empty(),
+  gasLimits?: Gas,
 ): Promise<Tx> {
   assert(
     setupCallRequests.length > 0 || appCallRequests.length > 0 || teardownCallRequest !== undefined,
     "Can't create public tx with no enqueued calls",
   );
   // use max limits
-  const gasLimits = new Gas(DEFAULT_DA_GAS_LIMIT, DEFAULT_L2_GAS_LIMIT);
+  gasLimits = gasLimits ?? new Gas(DEFAULT_DA_GAS_LIMIT, DEFAULT_L2_GAS_LIMIT);
 
   const forPublic = PartialPrivateTailPublicInputsForPublic.empty();
 
@@ -132,16 +133,15 @@ export async function createTxForPublicCalls(
     : Gas.empty();
   const gasSettings = new GasSettings(gasLimits, teardownGasLimits, maxFeesPerGas, GasFees.empty());
   const txContext = new TxContext(Fr.zero(), Fr.zero(), gasSettings);
-  const header = BlockHeader.empty();
-  header.globalVariables = globals;
+  const header = BlockHeader.empty({ globalVariables: globals });
   const constantData = new TxConstantData(header, txContext, Fr.zero(), Fr.zero());
-  const includeByTimestamp = 0n; // Not used in the simulator.
+  const expirationTimestamp = 0n; // Not used in the simulator.
 
   const txData = new PrivateKernelTailCircuitPublicInputs(
     constantData,
     /*gasUsed=*/ gasUsedByPrivate,
     feePayer,
-    includeByTimestamp,
+    expirationTimestamp,
     forPublic,
   );
 
@@ -172,13 +172,13 @@ export async function createTxForPrivateOnly(
   const gasSettings = new GasSettings(gasLimits, Gas.empty(), maxFeesPerGas, GasFees.empty());
   const txContext = new TxContext(Fr.zero(), Fr.zero(), gasSettings);
   const constantData = new TxConstantData(BlockHeader.empty(), txContext, Fr.zero(), Fr.zero());
-  const includeByTimestamp = 0n; // Not used in the simulator.
+  const expirationTimestamp = 0n; // Not used in the simulator.
 
   const txData = new PrivateKernelTailCircuitPublicInputs(
     constantData,
     /*gasUsed=*/ gasUsedByPrivate,
     feePayer,
-    includeByTimestamp,
+    expirationTimestamp,
     /*forPublic=*/ undefined,
     forRollup,
   );