@aztec/simulator 0.0.0-test.0 → 0.0.1-fake-c83136db25

package/src/public/avm/test_utils.ts

@@ -1,81 +1,64 @@
 import { Fr } from '@aztec/foundation/fields';
-import {
-  type ContractClassPublic,
-  type ContractInstanceWithAddress,
-  computePublicBytecodeCommitment,
-} from '@aztec/stdlib/contract';
+import type { ContractClassPublic, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 
 import type { jest } from '@jest/globals';
 import { mock } from 'jest-mock-extended';
 
-import type { WorldStateDB } from '../../public/public_db_sources.js';
+import type { PublicContractsDB, PublicTreesDB } from '../public_db_sources.js';
 import type { PublicSideEffectTraceInterface } from '../side_effect_trace_interface.js';
 
-export async function mockGetBytecode(worldStateDB: WorldStateDB, bytecode: Buffer) {
-  const commitment = await computePublicBytecodeCommitment(bytecode);
-  (worldStateDB as jest.Mocked<WorldStateDB>).getBytecodeCommitment.mockResolvedValue(commitment);
-}
-
 export function mockTraceFork(trace: PublicSideEffectTraceInterface, nestedTrace?: PublicSideEffectTraceInterface) {
   (trace as jest.Mocked<PublicSideEffectTraceInterface>).fork.mockReturnValue(
     nestedTrace ?? mock<PublicSideEffectTraceInterface>(),
   );
 }
 
-export function mockStorageRead(worldStateDB: WorldStateDB, value: Fr) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).storageRead.mockResolvedValue(value);
+export function mockStorageRead(worldStateDB: PublicTreesDB, value: Fr) {
+  (worldStateDB as jest.Mocked<PublicTreesDB>).storageRead.mockResolvedValue(value);
 }
 
 export function mockNoteHashCount(mockedTrace: PublicSideEffectTraceInterface, count: number) {
   (mockedTrace as jest.Mocked<PublicSideEffectTraceInterface>).getNoteHashCount.mockReturnValue(count);
 }
 
-export function mockStorageReadWithMap(worldStateDB: WorldStateDB, mockedStorage: Map<bigint, Fr>) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).storageRead.mockImplementation((_address, slot) =>
+export function mockStorageReadWithMap(worldStateDB: PublicTreesDB, mockedStorage: Map<bigint, Fr>) {
+  (worldStateDB as jest.Mocked<PublicTreesDB>).storageRead.mockImplementation((_address, slot) =>
     Promise.resolve(mockedStorage.get(slot.toBigInt()) ?? Fr.ZERO),
   );
 }
 
-export function mockGetBytecodeCommitment(worldStateDB: WorldStateDB, commitment: Fr) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).getBytecodeCommitment.mockResolvedValue(commitment);
-}
-
-export function mockNoteHashExists(worldStateDB: WorldStateDB, _leafIndex: Fr, value?: Fr) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).getCommitmentValue.mockImplementation((index: bigint) => {
-    if (index == _leafIndex.toBigInt()) {
+export function mockGetNoteHash(worldStateDB: PublicTreesDB, _leafIndex: bigint, value?: Fr) {
+  (worldStateDB as jest.Mocked<PublicTreesDB>).getNoteHash.mockImplementation((index: bigint) => {
+    if (index == _leafIndex && value) {
       return Promise.resolve(value);
     } else {
-      // This is ok for now since the traceing functions handle it
-      return Promise.resolve(undefined);
+      return Promise.resolve(Fr.ZERO);
     }
   });
 }
 
-export function mockGetNullifierIndex(worldStateDB: WorldStateDB, leafIndex: Fr, _ignoredValue?: Fr) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).getNullifierIndex.mockResolvedValue(leafIndex.toBigInt());
+export function mockCheckNullifierExists(worldStateDB: PublicTreesDB, exists: boolean, _ignoredValue?: Fr) {
+  (worldStateDB as jest.Mocked<PublicTreesDB>).checkNullifierExists.mockResolvedValue(exists);
 }
 
-export function mockL1ToL2MessageExists(
-  worldStateDB: WorldStateDB,
-  leafIndex: Fr,
-  value: Fr,
-  valueAtOtherIndices?: Fr,
-) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).getL1ToL2LeafValue.mockImplementation((index: bigint) => {
-    if (index == leafIndex.toBigInt()) {
+export function mockGetL1ToL2LeafValue(worldStateDB: PublicTreesDB, leafIndex: bigint, value?: Fr) {
+  (worldStateDB as jest.Mocked<PublicTreesDB>).getL1ToL2LeafValue.mockImplementation((index: bigint) => {
+    if (index == leafIndex && value) {
       return Promise.resolve(value);
     } else {
-      // any indices other than mockAtLeafIndex will return a different value
-      // (or undefined if no value is specified for other indices)
-      return Promise.resolve(valueAtOtherIndices!);
+      return Promise.resolve(Fr.ZERO!);
     }
   });
 }
 
-export function mockGetContractInstance(worldStateDB: WorldStateDB, contractInstance: ContractInstanceWithAddress) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).getContractInstance.mockResolvedValue(contractInstance);
+export function mockGetContractInstance(contractsDB: PublicContractsDB, contractInstance: ContractInstanceWithAddress) {
+  (contractsDB as jest.Mocked<PublicContractsDB>).getContractInstance.mockResolvedValue(contractInstance);
+}
+
+export function mockGetContractClass(contractsDB: PublicContractsDB, contractClass: ContractClassPublic) {
+  (contractsDB as jest.Mocked<PublicContractsDB>).getContractClass.mockResolvedValue(contractClass);
 }
 
-export function mockGetContractClass(worldStateDB: WorldStateDB, contractClass: ContractClassPublic) {
-  (worldStateDB as jest.Mocked<WorldStateDB>).getContractClass.mockResolvedValue(contractClass);
+export function mockGetBytecodeCommitment(contractsDB: PublicContractsDB, commitment: Fr) {
+  (contractsDB as jest.Mocked<PublicContractsDB>).getBytecodeCommitment.mockResolvedValue(commitment);
 }

package/src/public/db_interfaces.ts

@@ -0,0 +1,61 @@
+import type { Fr } from '@aztec/foundation/fields';
+import type { FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractClassPublic, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
+import type { UInt64 } from '@aztec/stdlib/types';
+
+/**
+ * Database interface for providing access to public state.
+ */
+export interface PublicStateDBInterface {
+  /**
+   * Reads a value from public storage, returning zero if none.
+   * @param contract - Owner of the storage.
+   * @param slot - Slot to read in the contract storage.
+   * @returns The current value in the storage slot.
+   */
+  storageRead(contract: AztecAddress, slot: Fr): Promise<Fr>;
+
+  /**
+   * Records a write to public storage.
+   * @param contract - Owner of the storage.
+   * @param slot - Slot to read in the contract storage.
+   * @param newValue - The new value to store.
+   */
+  storageWrite(contract: AztecAddress, slot: Fr, newValue: Fr): Promise<void>;
+}
+
+/**
+ * Database interface for providing access to public contract data.
+ */
+export interface PublicContractsDBInterface {
+  /**
+   * Returns a contract instance given its address and the given timestamp, or undefined if not exists.
+   * @param address - Address of the contract.
+   * @param timestamp - Timestamp to get the contract instance at. Contract updates might change the instance.
+   * @returns The contract instance or undefined if not found.
+   */
+  getContractInstance(address: AztecAddress, timestamp: UInt64): Promise<ContractInstanceWithAddress | undefined>;
+
+  /**
+   * Returns a publicly deployed contract class.
+   * @param contractClassId - ID of the contract class.
+   * @returns The contract class or undefined if not found
+   */
+  getContractClass(contractClassId: Fr): Promise<ContractClassPublic | undefined>;
+
+  /**
+   * Returns the commitment to the bytecode of a contract class.
+   * @param contractClassId - ID of the contract class.
+   * @returns The commitment to the bytecode or undefined if not found.
+   */
+  getBytecodeCommitment(contractClassId: Fr): Promise<Fr | undefined>;
+
+  /**
+   * Returns the function name of a contract's function given its selector.
+   * @param contractAddress - Address of the contract.
+   * @param selector - Selector of the function.
+   * @returns The name of the function or undefined if not found.
+   */
+  getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector): Promise<string | undefined>;
+}

package/src/{common → public}/debug_fn_name.ts

@@ -2,17 +2,17 @@ import type { Fr } from '@aztec/foundation/fields';
 import { FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 
-import type { WorldStateDB } from '../public/public_db_sources.js';
+import type { PublicContractsDBInterface } from './db_interfaces.js';
 
 export async function getPublicFunctionDebugName(
-  db: WorldStateDB,
+  db: PublicContractsDBInterface,
   contractAddress: AztecAddress,
   calldata: Fr[],
 ): Promise<string> {
   // Public function is dispatched and therefore the target function is passed in the first argument.
-  const targetFunction =
-    calldata[0] !== undefined
-      ? await db.getDebugFunctionName(contractAddress, FunctionSelector.fromField(calldata[0]))
-      : `<calldata[0] undefined> (Contract Address: ${contractAddress})`;
-  return `${targetFunction}`;
+  if (!calldata[0]) {
+    return `<calldata[0] undefined> (Contract Address: ${contractAddress})`;
+  }
+  const selector = FunctionSelector.fromField(calldata[0]);
+  return (await db.getDebugFunctionName(contractAddress, selector)) ?? selector.toString();
 }

package/src/public/executor_metrics.ts

@@ -1,3 +1,4 @@
+import type { RevertCode } from '@aztec/stdlib/avm';
 import {
   Attributes,
   type Histogram,
@@ -8,11 +9,16 @@ import {
   ValueType,
 } from '@aztec/telemetry-client';
 
-export class ExecutorMetrics {
+import type { ExecutorMetricsInterface } from './executor_metrics_interface.js';
+
+export class ExecutorMetrics implements ExecutorMetricsInterface {
   public readonly tracer: Tracer;
   private fnCount: UpDownCounter;
   private fnDuration: Histogram;
   private manaPerSecond: Histogram;
+  private manaUsed: Histogram;
+  private totalInstructionsExecuted: Histogram;
+  private txHashing: Histogram;
   private privateEffectsInsertions: Histogram;
 
   constructor(client: TelemetryClient, name = 'PublicExecutor') {
@@ -35,20 +41,58 @@ export class ExecutorMetrics {
       valueType: ValueType.INT,
     });
 
-    this.privateEffectsInsertions = meter.createHistogram(Metrics.PUBLIC_EXECUTION_PRIVATE_EFFECTS_INSERTION, {
+    this.manaUsed = meter.createHistogram(Metrics.PUBLIC_EXECUTOR_SIMULATION_MANA_USED, {
+      description: 'Total mana used',
+      unit: 'mana',
+      valueType: ValueType.INT,
+    });
+
+    this.totalInstructionsExecuted = meter.createHistogram(Metrics.PUBLIC_EXECUTOR_SIMULATION_TOTAL_INSTRUCTIONS, {
+      description: 'Total number of instructions executed',
+      unit: '#instructions',
+      valueType: ValueType.INT,
+    });
+
+    this.txHashing = meter.createHistogram(Metrics.PUBLIC_EXECUTOR_TX_HASHING, {
+      description: 'Tx hashing time',
+      unit: 'ms',
+      valueType: ValueType.INT,
+    });
+
+    this.privateEffectsInsertions = meter.createHistogram(Metrics.PUBLIC_EXECUTOR_PRIVATE_EFFECTS_INSERTION, {
       description: 'Private effects insertion time',
       unit: 'us',
       valueType: ValueType.INT,
     });
   }
 
-  recordFunctionSimulation(durationMs: number, manaUsed: number, fnName: string) {
+  startRecordingTxSimulation(_txLabel: string) {
+    // do nothing (unimplemented)
+  }
+
+  stopRecordingTxSimulation(_txLabel: string, _revertedCode?: RevertCode) {
+    // do nothing (unimplemented)
+  }
+
+  recordEnqueuedCallSimulation(
+    fnName: string,
+    durationMs: number,
+    manaUsed: number,
+    totalInstructionsExecuted: number,
+  ) {
     this.fnCount.add(1, {
       [Attributes.OK]: true,
       [Attributes.APP_CIRCUIT_NAME]: fnName,
-      [Attributes.MANA_USED]: manaUsed,
     });
-    this.fnDuration.record(Math.ceil(durationMs));
+    this.manaUsed.record(Math.ceil(manaUsed), {
+      [Attributes.APP_CIRCUIT_NAME]: fnName,
+    });
+    this.totalInstructionsExecuted.record(Math.ceil(totalInstructionsExecuted), {
+      [Attributes.APP_CIRCUIT_NAME]: fnName,
+    });
+    this.fnDuration.record(Math.ceil(durationMs), {
+      [Attributes.APP_CIRCUIT_NAME]: fnName,
+    });
     if (durationMs > 0 && manaUsed > 0) {
       const manaPerSecond = Math.round((manaUsed * 1000) / durationMs);
       this.manaPerSecond.record(manaPerSecond, {
@@ -57,7 +101,12 @@ export class ExecutorMetrics {
     }
   }
 
-  recordFunctionSimulationFailure() {
+  recordEnqueuedCallSimulationFailure(
+    _fnName: string,
+    _durationMs: number,
+    _manaUsed: number,
+    _totalInstructionsExecuted: number,
+  ) {
     this.fnCount.add(1, {
       [Attributes.OK]: false,
     });

package/src/public/executor_metrics_interface.ts

@@ -0,0 +1,14 @@
+import type { RevertCode } from '@aztec/stdlib/avm';
+
+export interface ExecutorMetricsInterface {
+  startRecordingTxSimulation(txLabel: string): void;
+  stopRecordingTxSimulation(txLabel: string, revertedCode?: RevertCode): void;
+  recordEnqueuedCallSimulation(fnName: string, durationMs: number, manaUsed: number, totalInstructions: number): void;
+  recordEnqueuedCallSimulationFailure(
+    fnName: string,
+    durationMs: number,
+    manaUsed: number,
+    totalInstructionsExecuted: number,
+  ): void;
+  recordPrivateEffectsInsertion(durationUs: number, type: 'revertible' | 'non-revertible'): void;
+}

package/src/public/fixtures/amm_test.ts

@@ -0,0 +1,331 @@
+import { GeneratorIndex } from '@aztec/constants';
+import { poseidon2HashWithSeparator } from '@aztec/foundation/crypto';
+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';
+import { setUpToken } from './token_test.js';
+
+const INITIAL_TOKEN_BALANCE = 1_000_000_000n;
+/**
+ * THIS TEST IS BRITTLE! If it breaks, don't try fixing it.
+ * `.skip` it or literally just delete it and notify AVM team.
+ * You do NOT need permission to remove this test!
+ */
+export async function ammTest(
+  tester: PublicTxSimulationTester,
+  logger: Logger,
+  tokenArtifact: ContractArtifact,
+  ammArtifact: ContractArtifact,
+  expectToBeTrue: (x: boolean) => void,
+) {
+  const timer = new Timer();
+
+  const admin = AztecAddress.fromNumber(42);
+  const sender = AztecAddress.fromNumber(111);
+
+  logger.debug(`Deploying tokens`);
+  const token0 = await setUpToken(tester, tokenArtifact, admin, expectToBeTrue, /*seed=*/ 0);
+  const token1 = await setUpToken(tester, tokenArtifact, admin, expectToBeTrue, /*seed=*/ 1);
+  const liquidityToken = await setUpToken(tester, tokenArtifact, admin, expectToBeTrue, /*seed=*/ 2);
+  logger.debug(`Deploying AMM`);
+  const constructorArgs = [token0, token1, liquidityToken];
+  const amm = await tester.registerAndDeployContract(
+    constructorArgs,
+    /*deployer=*/ admin,
+    ammArtifact,
+    /*skipNullifierInsertion=*/ false,
+    /*seed=*/ 3,
+  );
+
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  const ammConstructorResult = await tester.executeTxWithLabel(
+    /*txLabel=*/ 'AMM/constructor',
+    /*sender=*/ admin,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        fnName: 'constructor',
+        args: constructorArgs,
+        address: amm.address,
+      },
+    ],
+  );
+  expectToBeTrue(ammConstructorResult.revertCode.isOK());
+
+  logger.debug(`Setting AMM as minter for liquidity token`);
+
+  // set the AMM as the minter for the liquidity token
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  const setMinterResult = await tester.executeTxWithLabel(
+    /*txLabel=*/ 'AMM/set_minter',
+    /*sender=*/ admin,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      {
+        fnName: 'set_minter',
+        args: [/*minter=*/ amm, /*approve=*/ true],
+        address: liquidityToken.address,
+      },
+    ],
+  );
+  expectToBeTrue(setMinterResult.revertCode.isOK());
+
+  logger.debug(`Adding liquidity`);
+  const amount0Max = (INITIAL_TOKEN_BALANCE * 6n) / 10n;
+  const amount0Min = (INITIAL_TOKEN_BALANCE * 4n) / 10n;
+  const amount1Max = (INITIAL_TOKEN_BALANCE * 5n) / 10n;
+  const amount1Min = (INITIAL_TOKEN_BALANCE * 4n) / 10n;
+
+  const addLiquidityResult = await addLiquidity(
+    tester,
+    sender,
+    /*amm=*/ amm,
+    /*token0=*/ token0,
+    /*token1=*/ token1,
+    /*liquidityToken=*/ liquidityToken,
+    /*amount0Max=*/ amount0Max,
+    /*amount1Max=*/ amount1Max,
+    /*amount0Min=*/ amount0Min,
+    /*amount1Min=*/ amount1Min,
+  );
+  expectToBeTrue(addLiquidityResult.revertCode.isOK());
+
+  logger.debug(`Swapping tokens`);
+  const swapResult = await swapExactTokensForTokens(
+    tester,
+    sender,
+    /*amm=*/ amm,
+    /*tokenIn=*/ token0,
+    /*tokenOut=*/ token1,
+    /*amountIn=*/ amount0Min / 10n, // something smaller than total liquidity
+    /*amountOutMin=*/ amount1Min / 100n, // something even smaller
+  );
+  expectToBeTrue(swapResult.revertCode.isOK());
+
+  logger.debug(`Removing liquidity`);
+  const removeLiquidityResult = await removeLiquidity(
+    tester,
+    sender,
+    /*amm=*/ amm,
+    /*token0=*/ token0,
+    /*token1=*/ token1,
+    /*liquidityToken=*/ liquidityToken,
+    /*liquidity=*/ 100n,
+    /*amount0Min=*/ 1n, // remove some tiny amount
+    /*amount1Min=*/ 1n,
+  );
+  expectToBeTrue(removeLiquidityResult.revertCode.isOK());
+
+  logger.info(`AMM public tx simulator test took ${timer.ms()}ms\n`);
+}
+
+async function addLiquidity(
+  tester: PublicTxSimulationTester,
+  sender: AztecAddress,
+  amm: ContractInstanceWithAddress,
+  token0: ContractInstanceWithAddress,
+  token1: ContractInstanceWithAddress,
+  liquidityToken: ContractInstanceWithAddress,
+  amount0Max: bigint,
+  amount1Max: bigint,
+  amount0Min: bigint,
+  amount1Min: bigint,
+  _nonce?: bigint,
+) {
+  const refundToken0PartialNote = {
+    commitment: new Fr(42),
+  };
+  const refundToken1PartialNote = {
+    commitment: new Fr(66),
+  };
+  const liquidityPartialNote = {
+    commitment: new Fr(99),
+  };
+  const refundToken0PartialNoteValidityCommitment = await computePartialNoteValidityCommitment(
+    refundToken0PartialNote,
+    amm.address,
+  );
+  const refundToken1PartialNoteValidityCommitment = await computePartialNoteValidityCommitment(
+    refundToken1PartialNote,
+    amm.address,
+  );
+  const liquidityPartialNoteValidityCommitment = await computePartialNoteValidityCommitment(
+    liquidityPartialNote,
+    amm.address,
+  );
+
+  // We need to inject the validity commitments into the nullifier tree as that would be performed by the private token
+  // functions that are not invoked in this test.
+  await tester.insertNullifier(token0.address, refundToken0PartialNoteValidityCommitment);
+  await tester.insertNullifier(token1.address, refundToken1PartialNoteValidityCommitment);
+  await tester.insertNullifier(liquidityToken.address, liquidityPartialNoteValidityCommitment);
+
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  return await tester.executeTxWithLabel(
+    /*txLabel=*/ 'AMM/add_liquidity',
+    /*sender=*/ sender,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      // token0.transfer_to_public enqueues a call to _increase_public_balance
+      {
+        sender: token0.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_increase_public_balance',
+        args: [/*to=*/ amm.address, /*amount=*/ amount0Max],
+        address: token0.address,
+      },
+      // token1.transfer_to_public enqueues a call to _increase_public_balance
+      {
+        sender: token1.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_increase_public_balance',
+        args: [/*to=*/ amm.address, /*amount=*/ amount1Max],
+        address: token1.address,
+      },
+      // amm.add_liquidity enqueues a call to _add_liquidity
+      {
+        sender: amm.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_add_liquidity',
+        args: [
+          /*config=*/ {
+            token0: token0.address,
+            token1: token1.address,
+            // eslint-disable-next-line camelcase
+            liquidity_token: liquidityToken.address,
+          },
+          refundToken0PartialNote,
+          refundToken1PartialNote,
+          liquidityPartialNote,
+          amount0Max,
+          amount1Max,
+          amount0Min,
+          amount1Min,
+        ],
+        address: amm.address,
+      },
+    ],
+  );
+}
+
+async function swapExactTokensForTokens(
+  tester: PublicTxSimulationTester,
+  sender: AztecAddress,
+  amm: ContractInstanceWithAddress,
+  tokenIn: ContractInstanceWithAddress,
+  tokenOut: ContractInstanceWithAddress,
+  amountIn: bigint,
+  amountOutMin: bigint,
+  _nonce?: bigint,
+) {
+  const tokenOutPartialNote = {
+    commitment: new Fr(166),
+  };
+  const tokenOutPartialNoteValidityCommitment = await computePartialNoteValidityCommitment(
+    tokenOutPartialNote,
+    amm.address,
+  );
+
+  // We need to inject the validity commitment into the nullifier tree as that would be performed by the private token
+  // function that is not invoked in this test.
+  await tester.insertNullifier(tokenOut.address, tokenOutPartialNoteValidityCommitment);
+
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  return await tester.executeTxWithLabel(
+    /*txLabel=*/ 'AMM/swap_exact_tokens_for_tokens',
+    /*sender=*/ sender,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      // tokenIn.transfer_to_public enqueues a call to _increase_public_balance
+      {
+        sender: tokenIn.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_increase_public_balance',
+        args: [/*to=*/ amm.address, /*amount=*/ amountIn],
+        address: tokenIn.address,
+      },
+      {
+        sender: amm.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_swap_exact_tokens_for_tokens',
+        args: [tokenIn.address, tokenOut.address, amountIn, amountOutMin, tokenOutPartialNote],
+        address: amm.address,
+      },
+    ],
+  );
+}
+
+async function removeLiquidity(
+  tester: PublicTxSimulationTester,
+  sender: AztecAddress,
+  amm: ContractInstanceWithAddress,
+  token0: ContractInstanceWithAddress,
+  token1: ContractInstanceWithAddress,
+  liquidityToken: ContractInstanceWithAddress,
+  liquidity: bigint,
+  amount0Min: bigint,
+  amount1Min: bigint,
+  _nonce?: bigint,
+) {
+  const token0PartialNote = {
+    commitment: new Fr(111),
+  };
+  const token1PartialNote = {
+    commitment: new Fr(222),
+  };
+  const token0PartialNoteValidityCommitment = await computePartialNoteValidityCommitment(
+    token0PartialNote,
+    amm.address,
+  );
+  const token1PartialNoteValidityCommitment = await computePartialNoteValidityCommitment(
+    token1PartialNote,
+    amm.address,
+  );
+
+  // We need to inject the validity commitments into the nullifier tree as that would be performed by the private token
+  // functions that are not invoked in this test.
+  await tester.insertNullifier(token0.address, token0PartialNoteValidityCommitment);
+  await tester.insertNullifier(token1.address, token1PartialNoteValidityCommitment);
+
+  // EXECUTE! This means that if using AvmProvingTester subclass, it will PROVE the transaction!
+  return await tester.executeTxWithLabel(
+    /*txLabel=*/ 'AMM/remove_liquidity',
+    /*sender=*/ sender,
+    /*setupCalls=*/ [],
+    /*appCalls=*/ [
+      // liquidityToken.transfer_to_public enqueues a call to _increase_public_balance
+      {
+        sender: liquidityToken.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_increase_public_balance',
+        args: [/*to=*/ amm.address, /*amount=*/ liquidity],
+        address: liquidityToken.address,
+      },
+      // amm.remove_liquidity enqueues a call to _remove_liquidity
+      {
+        sender: amm.address, // INTERNAL FUNCTION! Sender must be 'this'.
+        fnName: '_remove_liquidity',
+        args: [
+          /*config=*/ {
+            token0: token0.address,
+            token1: token1.address,
+            // eslint-disable-next-line camelcase
+            liquidity_token: liquidityToken.address,
+          },
+          liquidity,
+          token0PartialNote,
+          token1PartialNote,
+          amount0Min,
+          amount1Min,
+        ],
+        address: amm.address,
+      },
+    ],
+  );
+}
+
+async function computePartialNoteValidityCommitment(partialNote: { commitment: Fr }, completer: AztecAddress) {
+  return await poseidon2HashWithSeparator(
+    [partialNote.commitment, completer],
+    GeneratorIndex.PARTIAL_NOTE_VALIDITY_COMMITMENT,
+  );
+}