@aztec/simulator 0.0.0-test.1 → 0.0.1-commit.023c3e5

package/src/public/fixtures/public_tx_simulation_tester.ts

@@ -1,53 +1,107 @@
-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 } 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 +111,29 @@ 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)] } },
   ): 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, DEFAULT_TEARDOWN_L2_GAS_LIMIT)
+        : Gas.empty(),
+      defaultGlobals(),
     );
   }
 
@@ -130,50 +144,145 @@ 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',
   ): Promise<PublicTxResult> {
-    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, firstNullifier);
+    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions);
 
     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;
 
-    const startTime = performance.now();
-    const avmResult = await simulator.simulate(tx);
-    const endTime = performance.now();
-    this.logger.debug(`Public transaction simulation took ${endTime - startTime}ms`);
+    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);
+
+    // 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,
+  ): Promise<PublicTxResult> {
+    return await this.simulateTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions, txLabel);
+  }
+
+  /**
+   * 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,
+  ): Promise<PublicTxResult> {
+    return await this.simulateTxWithLabel(
+      txLabel,
+      sender,
+      setupCalls,
+      appCalls,
+      teardownCall,
+      feePayer,
+      privateInsertions,
+    );
+  }
+
+  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);
+  }
 }
 
-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();
   }
 

package/src/public/fixtures/token_test.ts

@@ -0,0 +1,148 @@
+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';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+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();
+
+  const admin = AztecAddress.fromNumber(42);
+  const sender = AztecAddress.fromNumber(111);
+  const receiver = AztecAddress.fromNumber(222);
+
+  const token = await setUpToken(tester, tokenArtifact, admin, expectToBeTrue);
+
+  const mintAmount = 100n;
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  const mintResult = await tester.executeTxWithLabel(
+    /*txLabel=*/ 'Token/mint_to_public',
+    /*sender=*/ admin,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        address: token.address,
+        fnName: 'mint_to_public',
+        args: [/*to=*/ sender, mintAmount],
+      },
+    ],
+  );
+  expectToBeTrue(mintResult.revertCode.isOK());
+  await checkBalance(tester, token, sender, sender, mintAmount, expectToBeTrue, options);
+
+  const authwitNonce = new Fr(0);
+  const transferAmount = 50n;
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  const transferResult = await tester.executeTxWithLabel(
+    /*txLabel=*/ 'Token/transfer_in_public',
+    /*sender=*/ sender,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        address: token.address,
+        fnName: 'transfer_in_public',
+        args: [/*from=*/ sender, /*to=*/ receiver, transferAmount, authwitNonce],
+      },
+    ],
+  );
+  expectToBeTrue(transferResult.revertCode.isOK());
+  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(
+    /*txLabel=*/ 'Token/burn_public',
+    /*sender=*/ receiver,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        address: token.address,
+        fnName: 'burn_public',
+        args: [/*from=*/ receiver, transferAmount, authwitNonce],
+      },
+    ],
+  );
+  expectToBeTrue(burnResult.revertCode.isOK());
+  await checkBalance(tester, token, sender, receiver, 0n, expectToBeTrue, options);
+
+  logger.info(`TokenContract test took ${timer.ms()}ms\n`);
+}
+
+export async function setUpToken(
+  tester: PublicTxSimulationTester,
+  tokenArtifact: ContractArtifact,
+  admin: AztecAddress,
+  expectToBeTrue: (x: boolean) => void,
+  seed = 0,
+) {
+  const constructorArgs = [admin, /*name=*/ 'Token', /*symbol=*/ 'TOK', /*decimals=*/ new Fr(18)];
+  const token = await tester.registerAndDeployContract(
+    constructorArgs,
+    /*deployer=*/ admin,
+    tokenArtifact,
+    /*skipNullifierInsertion=*/ false,
+    seed,
+  );
+
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  const result = await tester.executeTxWithLabel(
+    /*txLabel=*/ 'Token/constructor',
+    /*sender=*/ admin,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        fnName: 'constructor',
+        args: constructorArgs,
+        address: token.address,
+      },
+    ],
+  );
+  expectToBeTrue(result.revertCode.isOK());
+  return token;
+}
+
+async function checkBalance(
+  tester: PublicTxSimulationTester,
+  token: ContractInstanceWithAddress,
+  sender: AztecAddress,
+  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(
+    /*txLabel=*/ 'Token/balance_of_public',
+    sender,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        address: token.address,
+        fnName: 'balance_of_public',
+        args: [/*owner=*/ account],
+        isStaticCall: true,
+      },
+    ],
+  );
+  expectToBeTrue(balResult.revertCode.isOK());
+  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);
+  }
+}