@aztec/simulator 0.72.1 → 0.73.0

package/src/avm/avm_simulator.ts

@@ -28,20 +28,14 @@ type OpcodeTally = {
   count: number;
   gas: Gas;
 };
-type PcTally = {
-  opcode: string;
-  count: number;
-  gas: Gas;
-};
 
 export class AvmSimulator {
   private log: Logger;
   private bytecode: Buffer | undefined;
   private opcodeTallies: Map<string, OpcodeTally> = new Map();
-  private pcTallies: Map<number, PcTally> = new Map();
 
   private tallyPrintFunction = () => {};
-  private tallyInstructionFunction = (_a: number, _b: string, _c: Gas) => {};
+  private tallyInstructionFunction = (_b: string, _c: Gas) => {};
 
   // Test Purposes only: Logger will not have the proper function name. Use this constructor for testing purposes
   // only. Otherwise, use build() below.
@@ -145,7 +139,6 @@ export class AvmSimulator {
       while (!machineState.getHalted()) {
         const [instruction, bytesRead] = decodeInstructionFromBytecode(bytecode, machineState.pc, this.instructionSet);
         const instrStartGas = machineState.gasLeft; // Save gas before executing instruction (for profiling)
-        const instrPc = machineState.pc; // Save PC before executing instruction (for profiling)
 
         this.log.trace(
           `[PC:${machineState.pc}] [IC:${instrCounter++}] ${instruction.toString()} (gasLeft l2=${
@@ -168,7 +161,7 @@ export class AvmSimulator {
           l2Gas: instrStartGas.l2Gas - machineState.l2GasLeft,
           daGas: instrStartGas.daGas - machineState.daGasLeft,
         };
-        this.tallyInstructionFunction(instrPc, instruction.constructor.name, gasUsed);
+        this.tallyInstructionFunction(instruction.constructor.name, gasUsed);
 
         if (machineState.pc >= bytecode.length) {
           this.log.warn('Passed end of program');
@@ -239,18 +232,12 @@ export class AvmSimulator {
     );
   }
 
-  private tallyInstruction(pc: number, opcode: string, gasUsed: Gas) {
+  private tallyInstruction(opcode: string, gasUsed: Gas) {
     const opcodeTally = this.opcodeTallies.get(opcode) || ({ count: 0, gas: { l2Gas: 0, daGas: 0 } } as OpcodeTally);
     opcodeTally.count++;
     opcodeTally.gas.l2Gas += gasUsed.l2Gas;
     opcodeTally.gas.daGas += gasUsed.daGas;
     this.opcodeTallies.set(opcode, opcodeTally);
-
-    const pcTally = this.pcTallies.get(pc) || ({ opcode: opcode, count: 0, gas: { l2Gas: 0, daGas: 0 } } as PcTally);
-    pcTally.count++;
-    pcTally.gas.l2Gas += gasUsed.l2Gas;
-    pcTally.gas.daGas += gasUsed.daGas;
-    this.pcTallies.set(pc, pcTally);
   }
 
   private printOpcodeTallies() {
@@ -261,16 +248,5 @@ export class AvmSimulator {
       // NOTE: don't care to clutter the logs with DA gas for now
       this.log.debug(`${opcode} executed ${tally.count} times consuming a total of ${tally.gas.l2Gas} L2 gas`);
     }
-
-    this.log.debug(`Printing tallies per PC sorted by #times each PC was executed...`);
-    const sortedPcs = Array.from(this.pcTallies.entries())
-      .sort((a, b) => b[1].count - a[1].count)
-      .filter((_, i) => i < 20);
-    for (const [pc, tally] of sortedPcs) {
-      // NOTE: don't care to clutter the logs with DA gas for now
-      this.log.debug(
-        `PC:${pc} containing opcode ${tally.opcode} executed ${tally.count} times consuming a total of ${tally.gas.l2Gas} L2 gas`,
-      );
-    }
   }
 }

package/src/avm/avm_tree.ts

@@ -109,7 +109,7 @@ export class AvmEphemeralForest {
    */
   async getSiblingPath(treeId: MerkleTreeId, index: bigint): Promise<Fr[]> {
     const tree = this.treeMap.get(treeId)!;
-    let path = tree.getSiblingPath(index);
+    let path = await tree.getSiblingPath(index);
     if (path === undefined) {
       // We dont have the sibling path in our tree - we have to get it from the DB
       path = (await this.treeDb.getSiblingPath(treeId, index)).toFields();
@@ -119,7 +119,7 @@ export class AvmEphemeralForest {
         const siblingIndex = index ^ 1n;
         const node = tree.getNode(siblingIndex, tree.depth - i);
         if (node !== undefined) {
-          const nodeHash = tree.hashTree(node, i + 1);
+          const nodeHash = await tree.hashTree(node, i + 1);
           if (!nodeHash.equals(path[i])) {
             path[i] = nodeHash;
           }
@@ -138,17 +138,17 @@ export class AvmEphemeralForest {
    * @param newLeafPreimage - The preimage of the new leaf to be inserted.
    * @returns The sibling path of the new leaf (i.e. the insertion path)
    */
-  appendIndexedTree<ID extends IndexedTreeId, T extends IndexedTreeLeafPreimage>(
+  async appendIndexedTree<ID extends IndexedTreeId, T extends IndexedTreeLeafPreimage>(
     treeId: ID,
     lowLeafIndex: bigint,
     lowLeafPreimage: T,
     newLeafPreimage: T,
-  ): Fr[] {
+  ): Promise<Fr[]> {
     const tree = this.treeMap.get(treeId)!;
-    const newLeaf = this.hashPreimage(newLeafPreimage);
+    const newLeaf = await this.hashPreimage(newLeafPreimage);
     const insertIndex = tree.leafCount;
 
-    const lowLeaf = this.hashPreimage(lowLeafPreimage);
+    const lowLeaf = await this.hashPreimage(lowLeafPreimage);
     // Update the low nullifier hash
     this.setIndexedUpdates(treeId, lowLeafIndex, lowLeafPreimage);
     tree.updateLeaf(lowLeaf, lowLeafIndex);
@@ -156,7 +156,7 @@ export class AvmEphemeralForest {
     tree.appendLeaf(newLeaf);
     this.setIndexedUpdates(treeId, insertIndex, newLeafPreimage);
 
-    return tree.getSiblingPath(insertIndex)!;
+    return (await tree.getSiblingPath(insertIndex))!;
   }
 
   /**
@@ -186,7 +186,7 @@ export class AvmEphemeralForest {
       const existingPublicDataSiblingPath = siblingPath;
       updatedPreimage.value = newValue;
 
-      tree.updateLeaf(this.hashPreimage(updatedPreimage), lowLeafIndex);
+      tree.updateLeaf(await this.hashPreimage(updatedPreimage), lowLeafIndex);
       this.setIndexedUpdates(treeId, lowLeafIndex, updatedPreimage);
       this._updateSortedKeys(treeId, [updatedPreimage.slot], [lowLeafIndex]);
 
@@ -212,7 +212,7 @@ export class AvmEphemeralForest {
       new Fr(preimage.getNextKey()),
       preimage.getNextIndex(),
     );
-    const insertionPath = this.appendIndexedTree(treeId, lowLeafIndex, updatedLowLeaf, newPublicDataLeaf);
+    const insertionPath = await this.appendIndexedTree(treeId, lowLeafIndex, updatedLowLeaf, newPublicDataLeaf);
 
     // Even though the low leaf key is not updated, we still need to update the sorted keys in case we have
     // not seen the low leaf before
@@ -281,7 +281,7 @@ export class AvmEphemeralForest {
     updatedLowNullifier.nextIndex = insertionIndex;
 
     const newNullifierLeaf = new NullifierLeafPreimage(nullifier, preimage.nextNullifier, preimage.nextIndex);
-    const insertionPath = this.appendIndexedTree(treeId, index, updatedLowNullifier, newNullifierLeaf);
+    const insertionPath = await this.appendIndexedTree(treeId, index, updatedLowNullifier, newNullifierLeaf);
 
     // Even though the low nullifier key is not updated, we still need to update the sorted keys in case we have
     // not seen the low nullifier before
@@ -308,11 +308,11 @@ export class AvmEphemeralForest {
    * @param value - The note hash to be appended
    * @returns The insertion result which contains the insertion path
    */
-  appendNoteHash(noteHash: Fr): Fr[] {
+  async appendNoteHash(noteHash: Fr): Promise<Fr[]> {
     const tree = this.treeMap.get(MerkleTreeId.NOTE_HASH_TREE)!;
     tree.appendLeaf(noteHash);
     // We use leafCount - 1 here because we would have just appended a leaf
-    const insertionPath = tree.getSiblingPath(tree.leafCount - 1n);
+    const insertionPath = await tree.getSiblingPath(tree.leafCount - 1n);
     return insertionPath!;
   }
 
@@ -484,14 +484,15 @@ export class AvmEphemeralForest {
   /**
    * This hashes the preimage to a field element
    */
-  hashPreimage<T extends TreeLeafPreimage>(preimage: T): Fr {
+  hashPreimage<T extends TreeLeafPreimage>(preimage: T): Promise<Fr> {
     const input = preimage.toHashInputs().map(x => Fr.fromBuffer(x));
     return poseidon2Hash(input);
   }
 
-  getTreeSnapshot(id: MerkleTreeId): AppendOnlyTreeSnapshot {
+  async getTreeSnapshot(id: MerkleTreeId): Promise<AppendOnlyTreeSnapshot> {
     const tree = this.treeMap.get(id)!;
-    return new AppendOnlyTreeSnapshot(tree.getRoot(), Number(tree.leafCount));
+    const root = await tree.getRoot();
+    return new AppendOnlyTreeSnapshot(root, Number(tree.leafCount));
   }
 }
 
@@ -551,19 +552,10 @@ const Leaf = (value: Fr): Leaf => ({
  */
 export class EphemeralAvmTree {
   private tree: Tree;
-  private readonly zeroHashes: Fr[];
   public frontier: Fr[];
 
-  private constructor(public leafCount: bigint, public depth: number) {
-    let zeroHash = Fr.zero();
-    // Can probably cache this elsewhere
-    const zeroHashes = [];
-    for (let i = 0; i < this.depth; i++) {
-      zeroHashes.push(zeroHash);
-      zeroHash = poseidon2Hash([zeroHash, zeroHash]);
-    }
-    this.tree = Leaf(zeroHash);
-    this.zeroHashes = zeroHashes;
+  private constructor(private root: Leaf, public leafCount: bigint, public depth: number, private zeroHashes: Fr[]) {
+    this.tree = root;
     this.frontier = [];
   }
 
@@ -573,7 +565,14 @@ export class EphemeralAvmTree {
     treeDb: MerkleTreeReadOperations,
     merkleId: MerkleTreeId,
   ): Promise<EphemeralAvmTree> {
-    const tree = new EphemeralAvmTree(forkedLeafCount, depth);
+    let zeroHash = Fr.zero();
+    // Can probably cache this elsewhere
+    const zeroHashes = [];
+    for (let i = 0; i < depth; i++) {
+      zeroHashes.push(zeroHash);
+      zeroHash = await poseidon2Hash([zeroHash, zeroHash]);
+    }
+    const tree = new EphemeralAvmTree(Leaf(zeroHash), forkedLeafCount, depth, zeroHashes);
     await tree.initializeFrontier(treeDb, merkleId);
     return tree;
   }
@@ -604,10 +603,10 @@ export class EphemeralAvmTree {
    * @param index - The index of the leaf for which a sibling path should be returned.
    * @returns The sibling path of the leaf, can fail if the path is not found
    */
-  getSiblingPath(index: bigint): Fr[] | undefined {
+  async getSiblingPath(index: bigint): Promise<Fr[] | undefined> {
     const searchPath = this._derivePathLE(index);
     // Handle cases where we error out
-    const { path, status } = this._getSiblingPath(searchPath, this.tree, []);
+    const { path, status } = await this._getSiblingPath(searchPath, this.tree, []);
     if (status === SiblingStatus.ERROR) {
       return undefined;
     }
@@ -695,7 +694,7 @@ export class EphemeralAvmTree {
   /**
    * Computes the root of the tree
    */
-  public getRoot(): Fr {
+  public getRoot(): Promise<Fr> {
     return this.hashTree(this.tree, this.depth);
   }
 
@@ -704,10 +703,13 @@ export class EphemeralAvmTree {
    * @param tree - The tree to be hashed
    * @param depth - The depth of the tree
    */
-  public hashTree(tree: Tree, depth: number): Fr {
+  public async hashTree(tree: Tree, depth: number): Promise<Fr> {
     switch (tree.tag) {
       case TreeType.NODE: {
-        return poseidon2Hash([this.hashTree(tree.leftTree, depth - 1), this.hashTree(tree.rightTree, depth - 1)]);
+        return poseidon2Hash([
+          await this.hashTree(tree.leftTree, depth - 1),
+          await this.hashTree(tree.rightTree, depth - 1),
+        ]);
       }
       case TreeType.LEAF: {
         return tree.value;
@@ -807,7 +809,7 @@ export class EphemeralAvmTree {
    * @param tree - The current tree
    * @param acc - The accumulated sibling path
    */
-  private _getSiblingPath(searchPath: number[], tree: Tree, acc: Fr[]): AccumulatedSiblingPath {
+  private async _getSiblingPath(searchPath: number[], tree: Tree, acc: Fr[]): Promise<AccumulatedSiblingPath> {
     // If we have reached the end of the path, we should be at a leaf or empty node
     // If it is a leaf, we check if the value is equal to the leaf value
     // If it is empty we check if the value is equal to zero
@@ -824,15 +826,15 @@ export class EphemeralAvmTree {
       case TreeType.NODE: {
         // Look at the next element of the path to decided if we go left or right, note this mutates!
         return searchPath.pop() === 0
-          ? this._getSiblingPath(
+          ? await this._getSiblingPath(
               searchPath,
               tree.leftTree,
-              [this.hashTree(tree.rightTree, searchPath.length)].concat(acc),
+              [await this.hashTree(tree.rightTree, searchPath.length)].concat(acc),
             )
-          : this._getSiblingPath(
+          : await this._getSiblingPath(
               searchPath,
               tree.rightTree,
-              [this.hashTree(tree.leftTree, searchPath.length)].concat(acc),
+              [await this.hashTree(tree.leftTree, searchPath.length)].concat(acc),
             );
       }
       // In these two situations we are exploring a subtree we dont have information about

package/src/avm/fixtures/avm_simulation_tester.ts

@@ -0,0 +1,105 @@
+import { type MerkleTreeWriteOperations } from '@aztec/circuit-types';
+import { GasFees, GlobalVariables } from '@aztec/circuits.js';
+import { encodeArguments } from '@aztec/foundation/abi';
+import { type AztecAddress } from '@aztec/foundation/aztec-address';
+import { Fr } from '@aztec/foundation/fields';
+import { openTmpStore } from '@aztec/kv-store/lmdb';
+import { MerkleTrees } from '@aztec/world-state';
+
+import { type AvmContractCallResult } from '../../avm/avm_contract_call_result.js';
+import {
+  getContractFunctionArtifact,
+  getFunctionSelector,
+  initContext,
+  initExecutionEnvironment,
+  resolveContractAssertionMessage,
+} from '../../avm/fixtures/index.js';
+import { WorldStateDB } from '../../public/public_db_sources.js';
+import { SideEffectTrace } from '../../public/side_effect_trace.js';
+import { AvmPersistableStateManager, AvmSimulator } from '../../server.js';
+import { BaseAvmSimulationTester } from './base_avm_simulation_tester.js';
+import { SimpleContractDataSource } from './simple_contract_data_source.js';
+
+const TIMESTAMP = new Fr(99833);
+const DEFAULT_GAS_FEES = new GasFees(2, 3);
+
+/**
+ * A test class that extends the BaseAvmSimulationTester to enable real-app testing of the core AvmSimulator.
+ * It provides an interface for simulating one top-level call at a time and maintains state between
+ * subsequent top-level calls.
+ */
+export class AvmSimulationTester extends BaseAvmSimulationTester {
+  constructor(
+    contractDataSource: SimpleContractDataSource,
+    merkleTrees: MerkleTreeWriteOperations,
+    skipContractDeployments = false,
+    private stateManager: AvmPersistableStateManager,
+  ) {
+    super(contractDataSource, merkleTrees, skipContractDeployments);
+  }
+
+  static async create(skipContractDeployments = false): Promise<AvmSimulationTester> {
+    const contractDataSource = new SimpleContractDataSource();
+    const merkleTrees = await (await MerkleTrees.new(openTmpStore())).fork();
+    const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
+    const trace = new SideEffectTrace();
+    const firstNullifier = new Fr(420000);
+    // FIXME: merkle ops should work, but I'm seeing frequent (but inconsistent) bytecode retrieval
+    // failures on 2nd call to simulateCall with merkle ops on
+    const stateManager = await AvmPersistableStateManager.create(
+      worldStateDB,
+      trace,
+      /*doMerkleOperations=*/ false,
+      firstNullifier,
+    );
+    return new AvmSimulationTester(contractDataSource, merkleTrees, skipContractDeployments, stateManager);
+  }
+
+  /**
+   * Simulate a top-level contract call.
+   */
+  async simulateCall(
+    sender: AztecAddress,
+    address: AztecAddress,
+    fnName: string,
+    args: any[],
+    isStaticCall = false,
+  ): Promise<AvmContractCallResult> {
+    const contractArtifact = await this.contractDataSource.getContractArtifact(address);
+    if (!contractArtifact) {
+      throw new Error(`Contract not found at address: ${address}`);
+    }
+    const fnSelector = await getFunctionSelector(fnName, contractArtifact);
+    const fnAbi = getContractFunctionArtifact(fnName, contractArtifact);
+    const encodedArgs = encodeArguments(fnAbi!, args);
+    const calldata = [fnSelector.toField(), ...encodedArgs];
+
+    const globals = GlobalVariables.empty();
+    globals.timestamp = TIMESTAMP;
+    globals.gasFees = DEFAULT_GAS_FEES;
+
+    const environment = initExecutionEnvironment({
+      calldata,
+      globals,
+      address,
+      sender,
+      isStaticCall,
+    });
+    const persistableState = this.stateManager.fork();
+    const context = initContext({ env: environment, persistableState });
+
+    // First we simulate (though it's not needed in this simple case).
+    const simulator = new AvmSimulator(context);
+    const result = await simulator.execute();
+    if (result.reverted) {
+      this.logger.error(`Error in ${fnName}:`);
+      this.logger.error(
+        resolveContractAssertionMessage(fnName, result.revertReason!, result.output, contractArtifact)!,
+      );
+    } else {
+      this.logger.info(`Simulation of function ${fnName} succeeded!`);
+      this.stateManager.merge(persistableState);
+    }
+    return result;
+  }
+}

package/src/avm/fixtures/base_avm_simulation_tester.ts

@@ -0,0 +1,104 @@
+import { MerkleTreeId, type MerkleTreeWriteOperations } from '@aztec/circuit-types';
+import {
+  type ContractClassPublic,
+  type ContractInstanceWithAddress,
+  DEPLOYER_CONTRACT_ADDRESS,
+  PUBLIC_DISPATCH_SELECTOR,
+  PublicDataWrite,
+  computeInitializationHash,
+} from '@aztec/circuits.js';
+import { computePublicDataTreeLeafSlot, siloNullifier } from '@aztec/circuits.js/hash';
+import { makeContractClassPublic, makeContractInstanceFromClassId } from '@aztec/circuits.js/testing';
+import { type ContractArtifact, FunctionSelector } from '@aztec/foundation/abi';
+import { AztecAddress } from '@aztec/foundation/aztec-address';
+import { type Fr } from '@aztec/foundation/fields';
+import { createLogger } from '@aztec/foundation/log';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+
+import { computeFeePayerBalanceStorageSlot } from '../../server.js';
+import { PUBLIC_DISPATCH_FN_NAME, getContractFunctionArtifact } from './index.js';
+import { type SimpleContractDataSource } from './simple_contract_data_source.js';
+
+/**
+ * An abstract test class that enables tests of real apps in the AVM without requiring e2e tests.
+ * It enables this by letting us (1) perform pseudo-contract-deployments (and registrations)
+ * that trigger merkle tree operations and (2) maintain a contract data source to store
+ * and retrieve contract classes and instances.
+ *
+ * This class is meant to be extended when writing tests for a specific simulation interface.
+ * For example, has been extended for testing of the core AvmSimulator, and again for the PublicTxSimulator,
+ * both of which benefit from such pseudo-deployments by populating merkle trees and a contract data source
+ * with contract information.
+ */
+export abstract class BaseAvmSimulationTester {
+  public logger = createLogger('avm-simulation-tester');
+
+  constructor(
+    public contractDataSource: SimpleContractDataSource,
+    public merkleTrees: MerkleTreeWriteOperations,
+    /* May want to skip contract deployment tree ops to test failed contract address nullifier checks on CALL */
+    private skipContractDeployments = false,
+  ) {}
+
+  async setFeePayerBalance(feePayer: AztecAddress, balance: Fr) {
+    const feeJuiceAddress = ProtocolContractAddress.FeeJuice;
+    const balanceSlot = await computeFeePayerBalanceStorageSlot(feePayer);
+    await this.setPublicStorage(feeJuiceAddress, balanceSlot, balance);
+  }
+
+  async setPublicStorage(address: AztecAddress, slot: Fr, value: Fr) {
+    const leafSlot = await computePublicDataTreeLeafSlot(address, slot);
+    // get existing preimage
+    const publicDataWrite = new PublicDataWrite(leafSlot, value);
+    await this.merkleTrees.batchInsert(MerkleTreeId.PUBLIC_DATA_TREE, [publicDataWrite.toBuffer()], 0);
+  }
+
+  /**
+   * Derive the contract class and instance with some seed.
+   * Add both to the contract data source along with the contract artifact.
+   */
+  async registerAndDeployContract(
+    constructorArgs: any[],
+    deployer: AztecAddress,
+    contractArtifact: ContractArtifact,
+    seed = 0,
+  ): Promise<ContractInstanceWithAddress> {
+    const bytecode = getContractFunctionArtifact(PUBLIC_DISPATCH_FN_NAME, contractArtifact)!.bytecode;
+    const contractClass = await makeContractClassPublic(
+      seed,
+      /*publicDispatchFunction=*/ { bytecode, selector: new FunctionSelector(PUBLIC_DISPATCH_SELECTOR) },
+    );
+
+    const constructorAbi = getContractFunctionArtifact('constructor', contractArtifact);
+    const initializationHash = await computeInitializationHash(constructorAbi, constructorArgs);
+    const contractInstance = await makeContractInstanceFromClassId(contractClass.id, seed, {
+      deployer,
+      initializationHash,
+    });
+
+    await this.addContractClass(contractClass, contractArtifact);
+    await this.addContractInstance(contractInstance);
+    return contractInstance;
+  }
+
+  getFirstContractInstance(): ContractInstanceWithAddress {
+    return this.contractDataSource.getFirstContractInstance();
+  }
+
+  addContractClass(contractClass: ContractClassPublic, contractArtifact: ContractArtifact): Promise<void> {
+    this.logger.debug(`Adding contract class with Id ${contractClass.id}`);
+    this.contractDataSource.addContractArtifact(contractClass.id, contractArtifact);
+    return this.contractDataSource.addContractClass(contractClass);
+  }
+
+  async addContractInstance(contractInstance: ContractInstanceWithAddress) {
+    if (!this.skipContractDeployments) {
+      const contractAddressNullifier = await siloNullifier(
+        AztecAddress.fromNumber(DEPLOYER_CONTRACT_ADDRESS),
+        contractInstance.address.toField(),
+      );
+      await this.merkleTrees.batchInsert(MerkleTreeId.NULLIFIER_TREE, [contractAddressNullifier.toBuffer()], 0);
+    }
+    await this.contractDataSource.addContractInstance(contractInstance);
+  }
+}

package/src/avm/fixtures/index.ts

@@ -1,6 +1,6 @@
 import { isNoirCallStackUnresolved } from '@aztec/circuit-types';
 import { GasFees, GlobalVariables, MAX_L2_GAS_PER_TX_PUBLIC_PORTION } from '@aztec/circuits.js';
-import { type FunctionArtifact, FunctionSelector } from '@aztec/foundation/abi';
+import { type ContractArtifact, type FunctionArtifact, FunctionSelector } from '@aztec/foundation/abi';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { EthAddress } from '@aztec/foundation/eth-address';
 import { Fr } from '@aztec/foundation/fields';
@@ -23,6 +23,8 @@ import { AvmPersistableStateManager } from '../journal/journal.js';
 import { NullifierManager } from '../journal/nullifiers.js';
 import { PublicStorage } from '../journal/public_storage.js';
 
+export const PUBLIC_DISPATCH_FN_NAME = 'public_dispatch';
+
 /**
  * Create a new AVM context with default values.
  */
@@ -124,15 +126,51 @@ export function randomMemoryFields(length: number): Field[] {
   return [...Array(length)].map(_ => new Field(Fr.random()));
 }
 
-export function getAvmTestContractFunctionSelector(functionName: string): FunctionSelector {
-  const artifact = AvmTestContractArtifact.functions.find(f => f.name === functionName)!;
-  assert(!!artifact, `Function ${functionName} not found in AvmTestContractArtifact`);
-  const params = artifact.parameters;
-  return FunctionSelector.fromNameAndParameters(artifact.name, params);
+export function getFunctionSelector(
+  functionName: string,
+  contractArtifact: ContractArtifact,
+): Promise<FunctionSelector> {
+  const fnArtifact = contractArtifact.functions.find(f => f.name === functionName)!;
+  assert(!!fnArtifact, `Function ${functionName} not found in ${contractArtifact.name}`);
+  const params = fnArtifact.parameters;
+  return FunctionSelector.fromNameAndParameters(fnArtifact.name, params);
+}
+
+export function getContractFunctionArtifact(
+  functionName: string,
+  contractArtifact: ContractArtifact,
+): FunctionArtifact | undefined {
+  const artifact = contractArtifact.functions.find(f => f.name === functionName)!;
+  if (!artifact) {
+    return undefined;
+  }
+  return artifact;
+}
+
+export function resolveContractAssertionMessage(
+  functionName: string,
+  revertReason: AvmRevertReason,
+  output: Fr[],
+  contractArtifact: ContractArtifact,
+): string | undefined {
+  traverseCauseChain(revertReason, cause => {
+    revertReason = cause as AvmRevertReason;
+  });
+
+  const functionArtifact = contractArtifact.functions.find(f => f.name === functionName);
+  if (!functionArtifact || !revertReason.noirCallStack || !isNoirCallStackUnresolved(revertReason.noirCallStack)) {
+    return undefined;
+  }
+
+  return resolveAssertionMessageFromRevertData(output, functionArtifact);
+}
+
+export function getAvmTestContractFunctionSelector(functionName: string): Promise<FunctionSelector> {
+  return getFunctionSelector(functionName, AvmTestContractArtifact);
 }
 
 export function getAvmTestContractArtifact(functionName: string): FunctionArtifact {
-  const artifact = AvmTestContractArtifact.functions.find(f => f.name === functionName)!;
+  const artifact = getContractFunctionArtifact(functionName, AvmTestContractArtifact);
   assert(
     !!artifact?.bytecode,
     `No bytecode found for function ${functionName}. Try re-running bootstrap.sh on the repository root.`,
@@ -150,14 +188,5 @@ export function resolveAvmTestContractAssertionMessage(
   revertReason: AvmRevertReason,
   output: Fr[],
 ): string | undefined {
-  traverseCauseChain(revertReason, cause => {
-    revertReason = cause as AvmRevertReason;
-  });
-
-  const functionArtifact = AvmTestContractArtifact.functions.find(f => f.name === functionName);
-  if (!functionArtifact || !revertReason.noirCallStack || !isNoirCallStackUnresolved(revertReason.noirCallStack)) {
-    return undefined;
-  }
-
-  return resolveAssertionMessageFromRevertData(output, functionArtifact);
+  return resolveContractAssertionMessage(functionName, revertReason, output, AvmTestContractArtifact);
 }