@aztec/simulator 0.0.0-test.0 → 0.0.1-commit.24de95ac

package/src/public/fixtures/minimal_public_tx.ts

@@ -0,0 +1,36 @@
+import { AvmCircuitInputs } from '@aztec/stdlib/avm';
+import { ProtocolContracts } from '@aztec/stdlib/tx';
+
+import avmMinimalCircuitInputsJson from '../../../artifacts/avm_minimal_inputs.json' with { type: 'json' };
+import { TypeTag } from '../avm/avm_memory_types.js';
+import { Add, Return, Set } from '../avm/opcodes/index.js';
+import { encodeToBytecode } from '../avm/serialization/bytecode_serialization.js';
+import { Opcode } from '../avm/serialization/instruction_serialization.js';
+import type { PublicTxResult } from '../public_tx_simulator/public_tx_simulator.js';
+import { testCustomBytecode } from './custom_bytecode_tester.js';
+import { PublicTxSimulationTester } from './public_tx_simulation_tester.js';
+
+export async function executeAvmMinimalPublicTx(tester: PublicTxSimulationTester): Promise<PublicTxResult> {
+  const minimalBytecode = encodeToBytecode([
+    new Set(/*indirect*/ 0, /*dstOffset*/ 0, TypeTag.UINT32, /*value*/ 1).as(Opcode.SET_8, Set.wireFormat8),
+    new Set(/*indirect*/ 0, /*dstOffset*/ 1, TypeTag.UINT32, /*value*/ 2).as(Opcode.SET_8, Set.wireFormat8),
+    new Add(/*indirect=*/ 0, /*aOffset=*/ 0, /*bOffset=*/ 1, /*dstOffset=*/ 2).as(Opcode.ADD_8, Add.wireFormat8),
+    new Return(/*indirect=*/ 0, /*copySizeOffset=*/ 0, /*returnOffset=*/ 2),
+  ]);
+
+  const result = await testCustomBytecode(minimalBytecode, tester, 'MinimalTx', 'AvmMinimalContract');
+
+  // Modify the protocol contracts to be all zeros
+  result.avmProvingRequest.inputs.hints.protocolContracts = ProtocolContracts.empty();
+  result.avmProvingRequest.inputs.publicInputs.protocolContracts = ProtocolContracts.empty();
+
+  return result;
+}
+
+/**
+ * Reads the AVM circuit inputs for the minimal public tx from a pre-generated JSON file.
+ * @returns The AvmCircuitInputs for the minimal public tx.
+ */
+export function readAvmMinimalPublicTxInputsFromFile(): AvmCircuitInputs {
+  return AvmCircuitInputs.schema.parse(avmMinimalCircuitInputsJson);
+}

package/src/public/fixtures/public_tx_simulation_tester.ts

@@ -1,25 +1,33 @@
-import { PUBLIC_DISPATCH_SELECTOR } from '@aztec/constants';
+import { asyncMap } from '@aztec/foundation/async-map';
 import { Fr } from '@aztec/foundation/fields';
-import { type ContractArtifact, FunctionSelector, encodeArguments } from '@aztec/stdlib/abi';
+import { type ContractArtifact, encodeArguments } from '@aztec/stdlib/abi';
 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 { MeasuredCppPublicTxSimulatorHintedDbs } from '../public_tx_simulator/cpp_public_tx_simulator.js';
+import { MeasuredPublicTxSimulator } from '../public_tx_simulator/measured_public_tx_simulator.js';
+import type { PublicTxResult } from '../public_tx_simulator/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;
   args: any[];
@@ -33,21 +41,43 @@ export type TestEnqueuedCall = {
  * 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(),
+    useCppSimulator: boolean = false,
   ) {
-    super(contractDataSource, merkleTrees);
+    super(contractDataSource, merkleTree);
+
+    const contractsDB = new PublicContractsDB(contractDataSource);
+    const config = {
+      doMerkleOperations: true,
+      skipFeeEnforcement: false,
+      clientInitiatedSimulation: true,
+    };
+    this.simulator = useCppSimulator
+      ? new MeasuredCppPublicTxSimulatorHintedDbs(merkleTree, contractsDB, globals, this.metrics, config)
+      : new MeasuredPublicTxSimulator(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,
+  ): 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();
+    return new PublicTxSimulationTester(merkleTree, contractDataSource, globals, metrics, useCppSimulator);
+  }
+
+  public setMetricsPrefix(prefix: string) {
+    this.metricsPrefix = prefix;
   }
 
   public async createTx(
@@ -57,69 +87,27 @@ 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,
-      );
-    }
-
-    return await createTxForPublicCalls(
-      firstNullifier,
-      setupExecutionRequests,
-      appExecutionRequests,
-      teardownExecutionRequest,
+    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;
+
+    this.txCount++;
+    return createTxForPublicCalls(
+      privateInsertions,
+      setupCallRequests,
+      appCallRequests,
+      teardownCallRequest,
       feePayer,
+      /*gasUsedByPrivate*/ Gas.empty(),
+      defaultGlobals(),
     );
   }
 
@@ -130,50 +118,107 @@ 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 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`);
+    // 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();
+  }
+
+  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}`);
+    }
+
+    const fnSelector = await getFunctionSelector(call.fnName, contractArtifact);
+    const fnAbi = getContractFunctionAbi(call.fnName, contractArtifact)!;
+    const encodedArgs = encodeArguments(fnAbi, call.args);
+    const 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 = DEFAULT_BLOCK_NUMBER;
   return globals;
 }

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

@@ -1,15 +1,10 @@
 import type { Fr } from '@aztec/foundation/fields';
 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 +22,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,21 +36,29 @@ 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> {
     throw new Error('Method not implemented.');
   }
@@ -82,14 +87,27 @@ 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> {
+    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()}`;
+    }
+    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();
+    }
+    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,140 @@
+import { Fr } from '@aztec/foundation/fields';
+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 async function tokenTest(
+  tester: PublicTxSimulationTester,
+  logger: Logger,
+  tokenArtifact: ContractArtifact,
+  expectToBeTrue: (x: boolean) => void,
+) {
+  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);
+
+  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, receiver, mintAmount - transferAmount, expectToBeTrue);
+  await checkBalance(tester, token, sender, receiver, transferAmount, expectToBeTrue);
+
+  // 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);
+
+  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,
+) {
+  // 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());
+  // should be 1 call with 1 return value that is expectedBalance
+  expectToBeTrue(balResult.processedPhases.length == 1);
+  expectToBeTrue(balResult.processedPhases[0].returnValues.length == 1);
+  expectToBeTrue(balResult.processedPhases[0].returnValues[0].values!.length == 1);
+  expectToBeTrue(balResult.processedPhases[0].returnValues[0].values![0].toBigInt() == expectedBalance);
+}