@aztec/simulator 0.0.1-commit.b655e406 → 0.0.1-commit.d3ec352c

package/src/public/avm/fixtures/utils.ts

@@ -1,4 +1,5 @@
 import { CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS } from '@aztec/constants';
+import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Fr } from '@aztec/foundation/fields';
 import {
   type ContractArtifact,
@@ -28,7 +29,7 @@ import type { AvmRevertReason } from '../errors.js';
 
 export const PUBLIC_DISPATCH_FN_NAME = 'public_dispatch';
 export const DEFAULT_TIMESTAMP: UInt64 = 99833n;
-export const DEFAULT_BLOCK_NUMBER = 42;
+export const DEFAULT_BLOCK_NUMBER = BlockNumber(42);
 
 /**
  * Create a new object with all the same properties as the original, except for the ones in the overrides object.

package/src/public/avm/opcodes/addressing_mode.ts

@@ -83,7 +83,7 @@ export class Addressing {
           baseAddr = mem.get(0);
           const baseAddrTag = baseAddr.getTag();
           if (!TaggedMemory.isValidMemoryAddressTag(baseAddrTag!)) {
-            throw TagCheckError.forOffset(0, TypeTag[baseAddrTag!], TypeTag[TypeTag.UINT32]);
+            throw TagCheckError.forBaseAddress(TypeTag[baseAddrTag!]);
           }
         }
         // Here we know that resolved[i] is at most 32 bits and baseAddr is at most 32 bits.
@@ -100,7 +100,7 @@ export class Addressing {
 
         // Final check.
         if (!TaggedMemory.isValidMemoryAddressTag(resolvedTag)) {
-          throw TagCheckError.forOffset(resolved[i], TypeTag[resolvedTag], TypeTag[TypeTag.UINT32]);
+          throw TagCheckError.forIndirectAddress(resolved[i], TypeTag[resolvedTag]);
         }
 
         resolved[i] = Number(resolvedValue.toBigInt());

package/src/public/avm/opcodes/misc.ts

@@ -44,10 +44,10 @@ export class DebugLog extends Instruction {
     const operands = [this.levelOffset, this.messageOffset, this.fieldsOffset, this.fieldsSizeOffset];
     const [levelOffset, messageOffset, fieldsOffset, fieldsSizeOffset] = addressing.resolve(operands, memory);
 
-    // DebugLog is a no-op except when doing client-initiated simulation.
+    // DebugLog is a no-op except unless the config is set to collect debug logs.
     // Note that we still do address resolution and basic tag-checking (above)
     // To avoid a special-case in the witness generator and circuit.
-    if (context.environment.clientInitiatedSimulation) {
+    if (context.environment.config.collectDebugLogs) {
       memory.checkTag(TypeTag.UINT8, levelOffset);
       const levelNumber = memory.get(levelOffset).toNumber();
       memory.checkTag(TypeTag.UINT32, fieldsSizeOffset);
@@ -56,11 +56,11 @@ export class DebugLog extends Instruction {
       const memoryReads = 1 /* level */ + 1 /* fieldsSize */ + this.messageSize /* message */ + fieldsSize; /* fields */
       if (
         context.persistableState.getDebugLogMemoryReads() + memoryReads >
-        context.environment.maxDebugLogMemoryReads
+        context.environment.config.collectionLimits.maxDebugLogMemoryReads
       ) {
         // Regular error on purpose: this is not a recoverable error.
         throw new Error(
-          `Max debug log memory reads exceeded: ${context.persistableState.getDebugLogMemoryReads() + memoryReads} > ${context.environment.maxDebugLogMemoryReads}`,
+          `Max debug log memory reads exceeded: ${context.persistableState.getDebugLogMemoryReads() + memoryReads} > ${context.environment.config.collectionLimits.maxDebugLogMemoryReads}`,
         );
       }
       context.persistableState.writeDebugLogMemoryReads(memoryReads);

package/src/public/avm/revert_reason.ts

@@ -18,13 +18,16 @@ async function createRevertReason(message: string, revertData: Fr[], context: Av
     message = context.machineState.collectedRevertInfo.recursiveRevertReason.message;
   }
 
-  const fnName = await context.persistableState.getPublicFunctionDebugName(context.environment);
+  const { functionSelector, functionName } = await context.persistableState.getPublicFunctionSelectorAndName(
+    context.environment,
+  );
 
   return new AvmRevertReason(
     message,
     /*failingFunction=*/ {
       contractAddress: context.environment.address,
-      functionName: fnName,
+      functionSelector,
+      functionName,
     },
     /*noirCallStack=*/ [...internalCallStack, context.machineState.pc].map(pc => `0.${pc}`),
     /*options=*/ { cause: nestedError },

package/src/public/avm/serialization/bytecode_serialization.ts

@@ -1,6 +1,13 @@
 import { type Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
 
-import { AvmExecutionError, AvmParsingError, InvalidOpcodeError, InvalidProgramCounterError } from '../errors.js';
+import {
+  AvmExecutionError,
+  AvmParsingError,
+  InvalidOpcodeError,
+  InvalidProgramCounterError,
+  InvalidTagValueError,
+  duringInstrFetch,
+} from '../errors.js';
 import {
   Add,
   And,
@@ -172,7 +179,7 @@ export function decodeInstructionFromBytecode(
   instructionSet: InstructionSet = INSTRUCTION_SET,
 ): [Instruction, number] {
   if (pc >= bytecode.length) {
-    throw new InvalidProgramCounterError(pc, bytecode.length);
+    throw new InvalidProgramCounterError(pc, bytecode.length - 1);
   }
 
   try {
@@ -182,7 +189,7 @@ export function decodeInstructionFromBytecode(
 
     if (opcode > MAX_OPCODE_VALUE) {
       throw new InvalidOpcodeError(
-        `Opcode ${opcode} (0x${opcode.toString(16)}) value is not in the range of valid opcodes.`,
+        `Opcode ${opcode} (0x${opcode.toString(16)}) value is not in the range of valid opcodes (at PC ${pc}).`,
       );
     }
 
@@ -192,13 +199,17 @@ export function decodeInstructionFromBytecode(
     }
 
     const instructionDeserializer: InstructionDeserializer = instructionDeserializerOrUndef;
+
     const instruction = instructionDeserializer(cursor);
     return [instruction, cursor.position() - startingPosition];
   } catch (error) {
-    if (error instanceof InvalidOpcodeError || error instanceof AvmExecutionError) {
-      throw error;
+    if (error instanceof InvalidTagValueError || error instanceof InvalidOpcodeError) {
+      throw duringInstrFetch(error, pc);
+    } else if (error instanceof AvmExecutionError) {
+      throw new AvmParsingError(`Instruction parsing error at pc ${pc}: ${error.message}`);
     } else {
-      throw new AvmParsingError(`${error}`);
+      const msg = error instanceof Error ? `: ${error.message}` : '';
+      throw new AvmParsingError(`Instruction fetching error at pc ${pc}${msg}`);
     }
   }
 }

package/src/public/contracts_db_checkpoint.ts

@@ -0,0 +1,41 @@
+import type { Fr } from '@aztec/foundation/fields';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractClassPublic, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
+
+export class ContractsDbCheckpoint {
+  private instances: Map<string, ContractInstanceWithAddress> = new Map();
+  private classes: Map<string, ContractClassPublic> = new Map();
+  private bytecodeCommitments: Map<string, Fr> = new Map();
+
+  public addInstance(address: AztecAddress, instance: ContractInstanceWithAddress): void {
+    this.instances.set(address.toString(), instance);
+  }
+
+  public addClass(classId: Fr, contractClass: ContractClassPublic): void {
+    this.classes.set(classId.toString(), contractClass);
+  }
+
+  public addBytecodeCommitment(classId: Fr, commitment: Fr): void {
+    this.bytecodeCommitments.set(classId.toString(), commitment);
+  }
+
+  public getInstance(address: AztecAddress): ContractInstanceWithAddress | undefined {
+    return this.instances.get(address.toString());
+  }
+
+  public getClass(classId: Fr): ContractClassPublic | undefined {
+    return this.classes.get(classId.toString());
+  }
+
+  public getBytecodeCommitment(classId: Fr): Fr | undefined {
+    return this.bytecodeCommitments.get(classId.toString());
+  }
+
+  public deepCopy(): ContractsDbCheckpoint {
+    const copy = new ContractsDbCheckpoint();
+    this.instances.forEach((value, key) => copy.instances.set(key, value));
+    this.classes.forEach((value, key) => copy.classes.set(key, value));
+    this.bytecodeCommitments.forEach((value, key) => copy.bytecodeCommitments.set(key, value));
+    return copy;
+  }
+}

package/src/public/db_interfaces.ts

@@ -58,4 +58,19 @@ export interface PublicContractsDBInterface {
    * @returns The name of the function or undefined if not found.
    */
   getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector): Promise<string | undefined>;
+
+  /**
+   * Creates a checkpoint of the current contract state.
+   */
+  createCheckpoint(): void;
+
+  /**
+   * Commits the current checkpoint, merging its state with the parent.
+   */
+  commitCheckpoint(): void;
+
+  /**
+   * Reverts the current checkpoint, discarding its state.
+   */
+  revertCheckpoint(): void;
 }

package/src/public/debug_fn_name.ts

@@ -16,3 +16,28 @@ export async function getPublicFunctionDebugName(
   const selector = FunctionSelector.fromField(calldata[0]);
   return (await db.getDebugFunctionName(contractAddress, selector)) ?? selector.toString();
 }
+
+/**
+ * Get the function selector and optional debug name for a public function.
+ * Returns the selector and name separately, with name only populated if a debug name is available.
+ * @param db - The contracts database
+ * @param contractAddress - The contract address
+ * @param calldata - The calldata (selector is in calldata[0])
+ * @returns An object with functionSelector (always if calldata[0] exists) and functionName (only if debug name found)
+ */
+export async function getPublicFunctionSelectorAndName(
+  db: PublicContractsDBInterface,
+  contractAddress: AztecAddress,
+  calldata: Fr[],
+): Promise<{ functionSelector?: FunctionSelector; functionName?: string }> {
+  // Public function is dispatched and therefore the target function is passed in the first argument.
+  if (!calldata[0]) {
+    return {};
+  }
+  const selector = FunctionSelector.fromField(calldata[0]);
+  const debugName = await db.getDebugFunctionName(contractAddress, selector);
+  return {
+    functionSelector: selector,
+    functionName: debugName ?? undefined,
+  };
+}

package/src/public/executor_metrics.ts

@@ -1,4 +1,5 @@
 import type { RevertCode } from '@aztec/stdlib/avm';
+import type { GasUsed } from '@aztec/stdlib/gas';
 import {
   Attributes,
   type Histogram,
@@ -70,7 +71,7 @@ export class ExecutorMetrics implements ExecutorMetricsInterface {
     // do nothing (unimplemented)
   }
 
-  stopRecordingTxSimulation(_txLabel: string, _revertedCode?: RevertCode) {
+  stopRecordingTxSimulation(_txLabel: string, _gasUsed?: GasUsed, _revertedCode?: RevertCode) {
     // do nothing (unimplemented)
   }
 

package/src/public/executor_metrics_interface.ts

@@ -1,8 +1,9 @@
 import type { RevertCode } from '@aztec/stdlib/avm';
+import type { GasUsed } from '@aztec/stdlib/gas';
 
 export interface ExecutorMetricsInterface {
   startRecordingTxSimulation(txLabel: string): void;
-  stopRecordingTxSimulation(txLabel: string, revertedCode?: RevertCode): void;
+  stopRecordingTxSimulation(txLabel: string, gasUsed?: GasUsed, revertedCode?: RevertCode): void;
   recordEnqueuedCallSimulation(fnName: string, durationMs: number, manaUsed: number, totalInstructions: number): void;
   recordEnqueuedCallSimulationFailure(
     fnName: string,

package/src/public/fixtures/custom_bytecode_tester.ts

@@ -1,7 +1,7 @@
 import { FunctionType, emptyContractArtifact, emptyFunctionArtifact } from '@aztec/stdlib/abi';
+import type { PublicTxResult } from '@aztec/stdlib/avm';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 
-import type { PublicTxResult } from '../public_tx_simulator/public_tx_simulator.js';
 import { PublicTxSimulationTester } from './public_tx_simulation_tester.js';
 
 /**

package/src/public/fixtures/minimal_public_tx.ts

@@ -1,4 +1,4 @@
-import { AvmCircuitInputs } from '@aztec/stdlib/avm';
+import { AvmCircuitInputs, type PublicTxResult } from '@aztec/stdlib/avm';
 import { ProtocolContracts } from '@aztec/stdlib/tx';
 
 import avmMinimalCircuitInputsJson from '../../../artifacts/avm_minimal_inputs.json' with { type: 'json' };
@@ -6,7 +6,6 @@ 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';
 
@@ -21,8 +20,8 @@ export async function executeAvmMinimalPublicTx(tester: PublicTxSimulationTester
   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();
+  result.hints!.protocolContracts = ProtocolContracts.empty();
+  result.publicInputs.protocolContracts = ProtocolContracts.empty();
 
   return result;
 }

package/src/public/fixtures/public_tx_simulation_tester.ts

@@ -1,6 +1,9 @@
+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/fields';
 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 { Gas, GasFees } from '@aztec/stdlib/gas';
 import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
@@ -16,9 +19,8 @@ import {
   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 { 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';
@@ -35,6 +37,25 @@ export type TestEnqueuedCall = {
   contractArtifact?: ContractArtifact;
 };
 
+const defaultConfig: PublicSimulatorConfig = PublicSimulatorConfig.from({
+  skipFeeEnforcement: false,
+  collectCallMetadata: true,
+  collectDebugLogs: true,
+  collectHints: 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
@@ -50,19 +71,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);
-    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);
+    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(
@@ -70,10 +89,14 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     globals: GlobalVariables = defaultGlobals(),
     metrics: TestExecutorMetrics = new TestExecutorMetrics(),
     useCppSimulator = false,
+    config: PublicSimulatorConfig = defaultConfig,
   ): Promise<PublicTxSimulationTester> {
     const contractDataSource = new SimpleContractDataSource();
     const merkleTree = await worldStateService.fork();
-    return new PublicTxSimulationTester(merkleTree, contractDataSource, globals, metrics, useCppSimulator);
+    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) {
@@ -106,7 +129,9 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       appCallRequests,
       teardownCallRequest,
       feePayer,
-      /*gasUsedByPrivate*/ Gas.empty(),
+      /*gasUsedByPrivate*/ teardownCall
+        ? new Gas(DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT)
+        : Gas.empty(),
       defaultGlobals(),
     );
   }
@@ -128,6 +153,11 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     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);
 
     // Something like this is often useful for debugging:
@@ -219,6 +249,6 @@ 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,3 +1,4 @@
+import type { BlockNumber } from '@aztec/foundation/branded-types';
 import type { Fr } from '@aztec/foundation/fields';
 import { createLogger } from '@aztec/foundation/log';
 import type { ContractArtifact, FunctionSelector } from '@aztec/stdlib/abi';
@@ -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

@@ -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,9 +139,10 @@ async function checkBalance(
     ],
   );
   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);
+  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);
+  }
 }