@aztec/simulator 0.72.1 → 0.74.0

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,104 @@
+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 { NativeWorldStateService } 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 NativeWorldStateService.tmp()).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,9 +1,10 @@
 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';
+import { AvmGadgetsTestContractArtifact } from '@aztec/noir-contracts.js/AvmGadgetsTest';
 import { AvmTestContractArtifact } from '@aztec/noir-contracts.js/AvmTest';
 
 import { strict as assert } from 'assert';
@@ -23,6 +24,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 +127,67 @@ 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`);
+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 getAvmGadgetsTestContractFunctionSelector(functionName: string): Promise<FunctionSelector> {
+  const artifact = AvmGadgetsTestContractArtifact.functions.find(f => f.name === functionName)!;
+  assert(!!artifact, `Function ${functionName} not found in AvmGadgetsTestContractArtifact`);
   const params = artifact.parameters;
   return FunctionSelector.fromNameAndParameters(artifact.name, params);
 }
 
 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.`,
+  );
+  return artifact;
+}
+
+export function getAvmGadgetsTestContractArtifact(functionName: string): FunctionArtifact {
+  const artifact = AvmGadgetsTestContractArtifact.functions.find(f => f.name === functionName)!;
   assert(
     !!artifact?.bytecode,
     `No bytecode found for function ${functionName}. Try re-running bootstrap.sh on the repository root.`,
@@ -145,16 +200,29 @@ export function getAvmTestContractBytecode(functionName: string): Buffer {
   return artifact.bytecode;
 }
 
+export function getAvmGadgetsTestContractBytecode(functionName: string): Buffer {
+  const artifact = getAvmGadgetsTestContractArtifact(functionName);
+  return artifact.bytecode;
+}
+
 export function resolveAvmTestContractAssertionMessage(
   functionName: string,
   revertReason: AvmRevertReason,
   output: Fr[],
+): string | undefined {
+  return resolveContractAssertionMessage(functionName, revertReason, output, AvmTestContractArtifact);
+}
+
+export function resolveAvmGadgetsTestContractAssertionMessage(
+  functionName: string,
+  revertReason: AvmRevertReason,
+  output: Fr[],
 ): string | undefined {
   traverseCauseChain(revertReason, cause => {
     revertReason = cause as AvmRevertReason;
   });
 
-  const functionArtifact = AvmTestContractArtifact.functions.find(f => f.name === functionName);
+  const functionArtifact = AvmGadgetsTestContractArtifact.functions.find(f => f.name === functionName);
   if (!functionArtifact || !revertReason.noirCallStack || !isNoirCallStackUnresolved(revertReason.noirCallStack)) {
     return undefined;
   }

package/src/avm/fixtures/simple_contract_data_source.ts

@@ -0,0 +1,98 @@
+import {
+  type ContractClassPublic,
+  type ContractDataSource,
+  type ContractInstanceWithAddress,
+  type FunctionSelector,
+  type PublicFunction,
+  computePublicBytecodeCommitment,
+} from '@aztec/circuits.js';
+import { type ContractArtifact } from '@aztec/foundation/abi';
+import { type AztecAddress } from '@aztec/foundation/aztec-address';
+import { type Fr } from '@aztec/foundation/fields';
+import { createLogger } from '@aztec/foundation/log';
+
+import { PUBLIC_DISPATCH_FN_NAME } from './index.js';
+
+/**
+ * This class is used during public/avm testing to function as a database of
+ * contract contract classes and instances. Tests can populate it with classes
+ * and instances and then probe it via the ContractDataSource interface.
+ *
+ * This class does not include any real merkle trees & merkle operations.
+ */
+export class SimpleContractDataSource implements ContractDataSource {
+  public logger = createLogger('simple-contract-data-source');
+
+  // maps contract class ID to class
+  private contractClasses: Map<string, ContractClassPublic> = new Map();
+  // maps contract instance address to instance
+  private contractInstances: Map<string, ContractInstanceWithAddress> = new Map();
+  // maps contract instance address to address
+  private contractArtifacts: Map<string, ContractArtifact> = new Map();
+
+  /////////////////////////////////////////////////////////////
+  // Helper functions not in the contract data source interface
+  getFirstContractInstance(): ContractInstanceWithAddress {
+    return this.contractInstances.values().next().value;
+  }
+
+  addContractArtifact(classId: Fr, artifact: ContractArtifact): void {
+    this.contractArtifacts.set(classId.toString(), artifact);
+  }
+
+  /////////////////////////////////////////////////////////////
+  // ContractDataSource function impelementations
+  getPublicFunction(_address: AztecAddress, _selector: FunctionSelector): Promise<PublicFunction> {
+    throw new Error('Method not implemented.');
+  }
+
+  getBlockNumber(): Promise<number> {
+    throw new Error('Method not implemented.');
+  }
+
+  getContractClass(id: Fr): Promise<ContractClassPublic | undefined> {
+    return Promise.resolve(this.contractClasses.get(id.toString()));
+  }
+
+  async getBytecodeCommitment(id: Fr): Promise<Fr | undefined> {
+    const contractClass = await this.getContractClass(id);
+    return Promise.resolve(computePublicBytecodeCommitment(contractClass!.packedBytecode));
+  }
+
+  getContract(address: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
+    return Promise.resolve(this.contractInstances.get(address.toString()));
+  }
+
+  getContractClassIds(): Promise<Fr[]> {
+    throw new Error('Method not implemented.');
+  }
+
+  async getContractArtifact(address: AztecAddress): Promise<ContractArtifact | undefined> {
+    const contractInstance = await this.getContract(address);
+    if (!contractInstance) {
+      this.logger.warn(`Contract not found at address: ${address}`);
+      return undefined;
+    }
+    this.logger.debug(`Retrieved contract artifact for address: ${address}`);
+    this.logger.debug(`Contract class ID: ${contractInstance.contractClassId}`);
+    return Promise.resolve(this.contractArtifacts.get(contractInstance!.contractClassId.toString()));
+  }
+
+  getContractFunctionName(_address: AztecAddress, _selector: FunctionSelector): Promise<string> {
+    return Promise.resolve(PUBLIC_DISPATCH_FN_NAME);
+  }
+
+  registerContractFunctionSignatures(_address: AztecAddress, _signatures: string[]): Promise<void> {
+    return Promise.resolve();
+  }
+
+  addContractClass(contractClass: ContractClassPublic): Promise<void> {
+    this.contractClasses.set(contractClass.id.toString(), contractClass);
+    return Promise.resolve();
+  }
+
+  addContractInstance(contractInstance: ContractInstanceWithAddress): Promise<void> {
+    this.contractInstances.set(contractInstance.address.toString(), contractInstance);
+    return Promise.resolve();
+  }
+}