@aztec/simulator 0.67.1 → 0.68.1

package/src/public/enqueued_call_side_effect_trace.ts

@@ -24,6 +24,7 @@ import {
   MAX_L2_TO_L1_MSGS_PER_TX,
   MAX_NOTE_HASHES_PER_TX,
   MAX_NULLIFIERS_PER_TX,
+  MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS,
   MAX_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
   MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
   MAX_UNENCRYPTED_LOGS_PER_TX,
@@ -42,7 +43,6 @@ import {
   PublicDataWrite,
   ScopedL2ToL1Message,
   ScopedLogHash,
-  type ScopedNoteHash,
   SerializableContractInstance,
   type TreeSnapshots,
 } from '@aztec/circuits.js';
@@ -59,6 +59,7 @@ import { type AvmExecutionEnvironment } from '../avm/avm_execution_environment.j
 import { type EnqueuedPublicCallExecutionResultWithSideEffects, type PublicFunctionCallResult } from './execution.js';
 import { SideEffectLimitReachedError } from './side_effect_errors.js';
 import { type PublicSideEffectTraceInterface } from './side_effect_trace_interface.js';
+import { UniqueClassIds } from './unique_class_ids.js';
 
 const emptyPublicDataPath = () => new Array(PUBLIC_DATA_TREE_HEIGHT).fill(Fr.zero());
 const emptyNoteHashPath = () => new Array(NOTE_HASH_TREE_HEIGHT).fill(Fr.zero());
@@ -74,7 +75,7 @@ export type SideEffects = {
   enqueuedCalls: PublicCallRequest[];
 
   publicDataWrites: PublicDataUpdateRequest[];
-  noteHashes: ScopedNoteHash[];
+  noteHashes: NoteHash[];
   nullifiers: Nullifier[];
   l2ToL1Msgs: ScopedL2ToL1Message[];
 
@@ -111,7 +112,7 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
   private publicDataWrites: PublicDataUpdateRequest[] = [];
   private protocolPublicDataWritesLength: number = 0;
   private userPublicDataWritesLength: number = 0;
-  private noteHashes: ScopedNoteHash[] = [];
+  private noteHashes: NoteHash[] = [];
   private nullifiers: Nullifier[] = [];
   private l2ToL1Messages: ScopedL2ToL1Message[] = [];
   private unencryptedLogs: UnencryptedL2Log[] = [];
@@ -129,6 +130,8 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
      *  otherwise the public kernel can fail to prove because TX limits are breached.
      */
     private readonly previousSideEffectArrayLengths: SideEffectArrayLengths = SideEffectArrayLengths.empty(),
+    /** We need to track the set of class IDs used for bytecode retrieval to deduplicate and enforce limits. */
+    private gotBytecodeFromClassIds: UniqueClassIds = new UniqueClassIds(),
   ) {
     this.log.debug(`Creating trace instance with startSideEffectCounter: ${startSideEffectCounter}`);
     this.sideEffectCounter = startSideEffectCounter;
@@ -146,6 +149,7 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
         this.previousSideEffectArrayLengths.l2ToL1Msgs + this.l2ToL1Messages.length,
         this.previousSideEffectArrayLengths.unencryptedLogs + this.unencryptedLogs.length,
       ),
+      this.gotBytecodeFromClassIds.fork(),
     );
   }
 
@@ -153,7 +157,7 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
     // sanity check to avoid merging the same forked trace twice
     assert(
       !forkedTrace.alreadyMergedIntoParent,
-      'Cannot merge a forked trace that has already been merged into its parent!',
+      'Bug! Cannot merge a forked trace that has already been merged into its parent!',
     );
     forkedTrace.alreadyMergedIntoParent = true;
 
@@ -172,10 +176,21 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
   }
 
   private mergeHints(forkedTrace: this) {
+    this.gotBytecodeFromClassIds.acceptAndMerge(forkedTrace.gotBytecodeFromClassIds);
+
     this.avmCircuitHints.enqueuedCalls.items.push(...forkedTrace.avmCircuitHints.enqueuedCalls.items);
 
     this.avmCircuitHints.contractInstances.items.push(...forkedTrace.avmCircuitHints.contractInstances.items);
-    this.avmCircuitHints.contractBytecodeHints.items.push(...forkedTrace.avmCircuitHints.contractBytecodeHints.items);
+
+    // merge in contract bytecode hints
+    // UniqueClassIds should prevent duplication
+    for (const [contractClassId, bytecodeHint] of forkedTrace.avmCircuitHints.contractBytecodeHints) {
+      assert(
+        !this.avmCircuitHints.contractBytecodeHints.has(contractClassId),
+        'Bug preventing duplication of contract bytecode hints',
+      );
+      this.avmCircuitHints.contractBytecodeHints.set(contractClassId, bytecodeHint);
+    }
 
     this.avmCircuitHints.publicDataReads.items.push(...forkedTrace.avmCircuitHints.publicDataReads.items);
     this.avmCircuitHints.publicDataWrites.items.push(...forkedTrace.avmCircuitHints.publicDataWrites.items);
@@ -194,6 +209,10 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
     this.sideEffectCounter++;
   }
 
+  public getNoteHashCount() {
+    return this.previousSideEffectArrayLengths.noteHashes + this.noteHashes.length;
+  }
+
   public tracePublicStorageRead(
     contractAddress: AztecAddress,
     slot: Fr,
@@ -272,19 +291,12 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
     // NOTE: counter does not increment for note hash checks (because it doesn't rely on pending note hashes)
   }
 
-  public traceNewNoteHash(
-    contractAddress: AztecAddress,
-    noteHash: Fr,
-    leafIndex: Fr = Fr.zero(),
-    path: Fr[] = emptyNoteHashPath(),
-  ) {
+  public traceNewNoteHash(noteHash: Fr, leafIndex: Fr = Fr.zero(), path: Fr[] = emptyNoteHashPath()) {
     if (this.noteHashes.length + this.previousSideEffectArrayLengths.noteHashes >= MAX_NOTE_HASHES_PER_TX) {
       throw new SideEffectLimitReachedError('note hash', MAX_NOTE_HASHES_PER_TX);
     }
 
-    // TODO(dbanks12): make unique and silo instead of scoping
-    //const siloedNoteHash = siloNoteHash(contractAddress, noteHash);
-    this.noteHashes.push(new NoteHash(noteHash, this.sideEffectCounter).scope(contractAddress));
+    this.noteHashes.push(new NoteHash(noteHash, this.sideEffectCounter));
     this.log.debug(`NEW_NOTE_HASH cnt: ${this.sideEffectCounter}`);
     this.avmCircuitHints.noteHashWrites.items.push(new AvmAppendTreeHint(leafIndex, noteHash, path));
     this.incrementSideEffectCounter();
@@ -409,6 +421,14 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
     lowLeafIndex: Fr = Fr.zero(),
     lowLeafPath: Fr[] = emptyNullifierPath(),
   ) {
+    // FIXME: The way we are hinting contract bytecodes is fundamentally broken.
+    // We are mapping contract class ID to a bytecode hint
+    // But a bytecode hint is tied to a contract INSTANCE.
+    // What if you encounter another contract instance with the same class ID?
+    // We can't hint that instance too since there is already an entry in the hints set that class ID.
+    // But without that instance hinted, the circuit can't prove that the called contract address
+    // actually corresponds to any class ID.
+
     const membershipHint = new AvmNullifierReadTreeHint(lowLeafPreimage, lowLeafIndex, lowLeafPath);
     const instance = new AvmContractInstanceHint(
       contractAddress,
@@ -420,13 +440,57 @@ export class PublicEnqueuedCallSideEffectTrace implements PublicSideEffectTraceI
       contractInstance.publicKeys,
       membershipHint,
     );
-    // We need to deduplicate the contract instances based on addresses
-    this.avmCircuitHints.contractBytecodeHints.items.push(
-      new AvmContractBytecodeHints(bytecode, instance, contractClass),
-    );
+
+    // Always hint the contract instance separately from the bytecode hint.
+    // Since the bytecode hints are keyed by class ID, we need to hint the instance separately
+    // since there might be multiple instances hinted for the same class ID.
+    this.avmCircuitHints.contractInstances.items.push(instance);
     this.log.debug(
-      `Bytecode retrieval for contract execution traced: exists=${exists}, instance=${jsonStringify(contractInstance)}`,
+      `Tracing contract instance for bytecode retrieval: exists=${exists}, instance=${jsonStringify(contractInstance)}`,
+    );
+
+    if (!exists) {
+      // this ensures there are no duplicates
+      this.log.debug(`Contract address ${contractAddress} does not exist. Not tracing bytecode & class ID.`);
+      return;
+    }
+    // We already hinted this bytecode. No need to
+    // Don't we still need to hint if the class ID already exists?
+    // Because the circuit needs to prove that the called contract address corresponds to the class ID.
+    // To do so, the circuit needs to know the class ID in the
+    if (this.gotBytecodeFromClassIds.has(contractInstance.contractClassId.toString())) {
+      // this ensures there are no duplicates
+      this.log.debug(
+        `Contract class id ${contractInstance.contractClassId.toString()} already exists in previous hints`,
+      );
+      return;
+    }
+
+    // If we could actually allow contract calls after the limit was reached, we would hint even if we have
+    // surpassed the limit of unique class IDs (still trace the failed bytecode retrieval)
+    // because the circuit needs to know the class ID to know when the limit is hit.
+    // BUT, the issue with this approach is that the sequencer could lie and say "this call was to a new class ID",
+    // and the circuit cannot prove that it's not true without deriving the class ID from bytecode,
+    // proving that it corresponds to the called contract address, and proving that the class ID wasn't already
+    // present/used. That would require more bytecode hashing which is exactly what this limit exists to avoid.
+    if (this.gotBytecodeFromClassIds.size() >= MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS) {
+      this.log.debug(
+        `Bytecode retrieval failure for contract class ID ${contractInstance.contractClassId.toString()} (limit reached)`,
+      );
+      throw new SideEffectLimitReachedError(
+        'contract calls to unique class IDs',
+        MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS,
+      );
+    }
+
+    this.log.debug(`Tracing bytecode & contract class for bytecode retrieval: class=${jsonStringify(contractClass)}`);
+    this.avmCircuitHints.contractBytecodeHints.set(
+      contractInstance.contractClassId.toString(),
+      new AvmContractBytecodeHints(bytecode, instance, contractClass),
     );
+    // After adding the bytecode hint, mark the classId as retrieved to avoid duplication.
+    // The above map alone isn't sufficient because we need to check the parent trace's (and its parent) as well.
+    this.gotBytecodeFromClassIds.add(contractInstance.contractClassId.toString());
   }
 
   /**

package/src/public/execution.ts

@@ -20,7 +20,6 @@ import {
   RevertCode,
   type ScopedL2ToL1Message,
   type ScopedLogHash,
-  type ScopedNoteHash,
   type TreeLeafReadRequest,
 } from '@aztec/circuits.js';
 import { computeVarArgsHash } from '@aztec/circuits.js/hash';
@@ -29,7 +28,7 @@ export interface PublicSideEffects {
   /** The contract storage update requests performed. */
   publicDataWrites: PublicDataUpdateRequest[];
   /** The new note hashes to be inserted into the note hashes tree. */
-  noteHashes: ScopedNoteHash[];
+  noteHashes: NoteHash[];
   /** The new nullifiers to be inserted into the nullifier tree. */
   nullifiers: Nullifier[];
   /** The new l2 to l1 messages generated to be inserted into the messages tree. */

package/src/public/executor_metrics.ts

@@ -13,6 +13,7 @@ export class ExecutorMetrics {
   private fnCount: UpDownCounter;
   private fnDuration: Histogram;
   private manaPerSecond: Histogram;
+  private privateEffectsInsertions: Histogram;
 
   constructor(client: TelemetryClient, name = 'PublicExecutor') {
     this.tracer = client.getTracer(name);
@@ -33,6 +34,12 @@ export class ExecutorMetrics {
       unit: 'mana/s',
       valueType: ValueType.INT,
     });
+
+    this.privateEffectsInsertions = meter.createHistogram(Metrics.PUBLIC_EXECUTION_PRIVATE_EFFECTS_INSERTION, {
+      description: 'Private effects insertion time',
+      unit: 'us',
+      valueType: ValueType.INT,
+    });
   }
 
   recordFunctionSimulation(durationMs: number, manaUsed: number, fnName: string) {
@@ -55,4 +62,10 @@ export class ExecutorMetrics {
       [Attributes.OK]: false,
     });
   }
+
+  recordPrivateEffectsInsertion(durationUs: number, type: 'revertible' | 'non-revertible') {
+    this.privateEffectsInsertions.record(Math.ceil(durationUs), {
+      [Attributes.REVERTIBILITY]: type,
+    });
+  }
 }

package/src/public/fee_payment.ts

@@ -1,13 +1,14 @@
 import { computePublicDataTreeLeafSlot, deriveStorageSlotInMap } from '@aztec/circuits.js/hash';
 import { type AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr } from '@aztec/foundation/fields';
-import { ProtocolContractAddress, ProtocolContractArtifact } from '@aztec/protocol-contracts';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { FeeJuiceArtifact } from '@aztec/protocol-contracts/fee-juice';
 
 /**
  * Computes the storage slot within the Fee Juice contract for the balance of the fee payer.
  */
 export function computeFeePayerBalanceStorageSlot(feePayer: AztecAddress) {
-  return deriveStorageSlotInMap(ProtocolContractArtifact.FeeJuice.storageLayout.balances.slot, feePayer);
+  return deriveStorageSlotInMap(FeeJuiceArtifact.storageLayout.balances.slot, feePayer);
 }
 
 /**

package/src/public/fixtures/index.ts

@@ -1,9 +1,10 @@
-import { MerkleTreeId, PublicExecutionRequest, Tx } from '@aztec/circuit-types';
+import { MerkleTreeId, type MerkleTreeWriteOperations, PublicExecutionRequest, Tx } from '@aztec/circuit-types';
 import {
   type AvmCircuitInputs,
   BlockHeader,
   CallContext,
   type ContractClassPublic,
+  type ContractDataSource,
   type ContractInstanceWithAddress,
   DEFAULT_GAS_LIMIT,
   DEPLOYER_CONTRACT_ADDRESS,
@@ -13,6 +14,7 @@ import {
   GasSettings,
   GlobalVariables,
   MAX_L2_GAS_PER_TX_PUBLIC_PORTION,
+  MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS,
   PartialPrivateTailPublicInputsForPublic,
   PrivateKernelTailCircuitPublicInputs,
   type PublicFunction,
@@ -29,60 +31,55 @@ import { type ContractArtifact, type FunctionArtifact } from '@aztec/foundation/
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { openTmpStore } from '@aztec/kv-store/lmdb';
-import { AvmTestContractArtifact } from '@aztec/noir-contracts.js';
-import { PublicTxSimulator, WorldStateDB } from '@aztec/simulator';
+import { AvmTestContractArtifact } from '@aztec/noir-contracts.js/AvmTest';
+import {
+  AvmEphemeralForest,
+  AvmSimulator,
+  PublicEnqueuedCallSideEffectTrace,
+  PublicTxSimulator,
+  WorldStateDB,
+} from '@aztec/simulator';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 import { MerkleTrees } from '@aztec/world-state';
 
 import { strict as assert } from 'assert';
 
+import { initContext, initExecutionEnvironment, initPersistableStateManager } from '../../avm/fixtures/index.js';
+
+const TIMESTAMP = new Fr(99833);
+
 export async function simulateAvmTestContractGenerateCircuitInputs(
   functionName: string,
-  calldata: Fr[] = [],
+  args: Fr[] = [],
   expectRevert: boolean = false,
+  contractDataSource = new MockedAvmTestContractDataSource(),
   assertionErrString?: string,
 ): Promise<AvmCircuitInputs> {
-  const sender = AztecAddress.random();
-  const functionSelector = getAvmTestContractFunctionSelector(functionName);
-  calldata = [functionSelector.toField(), ...calldata];
+  const globals = GlobalVariables.empty();
+  globals.timestamp = TIMESTAMP;
 
-  const globalVariables = GlobalVariables.empty();
-  globalVariables.gasFees = GasFees.empty();
-  globalVariables.timestamp = new Fr(99833);
-
-  const telemetry = new NoopTelemetryClient();
-  const merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork();
-  const contractDataSource = new MockedAvmTestContractDataSource();
+  const merkleTrees = await (await MerkleTrees.new(openTmpStore(), new NoopTelemetryClient())).fork();
+  await contractDataSource.deployContracts(merkleTrees);
   const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
 
-  const contractInstance = contractDataSource.contractInstance;
-  const contractAddressNullifier = siloNullifier(
-    AztecAddress.fromNumber(DEPLOYER_CONTRACT_ADDRESS),
-    contractInstance.address.toField(),
-  );
-  await merkleTrees.batchInsert(MerkleTreeId.NULLIFIER_TREE, [contractAddressNullifier.toBuffer()], 0);
-  // other contract address used by the bulk test's GETCONTRACTINSTANCE test
-  const otherContractAddressNullifier = siloNullifier(
-    AztecAddress.fromNumber(DEPLOYER_CONTRACT_ADDRESS),
-    contractDataSource.otherContractInstance.address.toField(),
-  );
-  await merkleTrees.batchInsert(MerkleTreeId.NULLIFIER_TREE, [otherContractAddressNullifier.toBuffer()], 0);
-
   const simulator = new PublicTxSimulator(
     merkleTrees,
     worldStateDB,
     new NoopTelemetryClient(),
-    globalVariables,
+    globals,
     /*doMerkleOperations=*/ true,
   );
 
+  const sender = AztecAddress.random();
+  const functionSelector = getAvmTestContractFunctionSelector(functionName);
+  args = [functionSelector.toField(), ...args];
   const callContext = new CallContext(
     sender,
-    contractInstance.address,
+    contractDataSource.firstContractInstance.address,
     contractDataSource.fnSelector,
     /*isStaticCall=*/ false,
   );
-  const executionRequest = new PublicExecutionRequest(callContext, calldata);
+  const executionRequest = new PublicExecutionRequest(callContext, args);
 
   const tx: Tx = createTxForPublicCall(executionRequest);
 
@@ -103,6 +100,46 @@ export async function simulateAvmTestContractGenerateCircuitInputs(
   return avmCircuitInputs;
 }
 
+export async function simulateAvmTestContractCall(
+  functionName: string,
+  args: Fr[] = [],
+  expectRevert: boolean = false,
+  contractDataSource = new MockedAvmTestContractDataSource(),
+) {
+  const globals = GlobalVariables.empty();
+  globals.timestamp = TIMESTAMP;
+
+  const merkleTrees = await (await MerkleTrees.new(openTmpStore(), new NoopTelemetryClient())).fork();
+  await contractDataSource.deployContracts(merkleTrees);
+  const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
+
+  const trace = new PublicEnqueuedCallSideEffectTrace();
+  const ephemeralTrees = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface());
+  const persistableState = initPersistableStateManager({
+    worldStateDB,
+    trace,
+    merkleTrees: ephemeralTrees,
+    doMerkleOperations: true,
+  });
+
+  const sender = AztecAddress.random();
+  const functionSelector = getAvmTestContractFunctionSelector(functionName);
+  args = [functionSelector.toField(), ...args];
+  const environment = initExecutionEnvironment({
+    calldata: args,
+    globals,
+    address: contractDataSource.firstContractInstance.address,
+    sender,
+  });
+  const context = initContext({ env: environment, persistableState });
+
+  // First we simulate (though it's not needed in this simple case).
+  const simulator = new AvmSimulator(context);
+  const results = await simulator.execute();
+
+  expect(results.reverted).toBe(expectRevert);
+}
+
 /**
  * Craft a carrier transaction for a public call for simulation by PublicTxSimulator.
  */
@@ -125,7 +162,7 @@ export function createTxForPublicCall(
   }
 
   const teardownGasLimits = isTeardown ? gasLimits : Gas.empty();
-  const gasSettings = new GasSettings(gasLimits, teardownGasLimits, GasFees.empty());
+  const gasSettings = new GasSettings(gasLimits, teardownGasLimits, GasFees.empty(), GasFees.empty());
   const txContext = new TxContext(Fr.zero(), Fr.zero(), gasSettings);
   const constantData = new TxConstantData(BlockHeader.empty(), txContext, Fr.zero(), Fr.zero());
 
@@ -144,23 +181,47 @@ export function createTxForPublicCall(
   return tx;
 }
 
-export class MockedAvmTestContractDataSource {
+export class MockedAvmTestContractDataSource implements ContractDataSource {
   private fnName = 'public_dispatch';
+  public fnSelector: FunctionSelector = getAvmTestContractFunctionSelector(this.fnName);
   private bytecode: Buffer;
-  public fnSelector: FunctionSelector;
   private publicFn: PublicFunction;
-  private contractClass: ContractClassPublic;
-  public contractInstance: ContractInstanceWithAddress;
   private bytecodeCommitment: Fr;
+
+  // maps contract class ID to class
+  private contractClasses: Map<string, ContractClassPublic> = new Map();
+  // maps contract instance address to instance
+  public contractInstances: Map<string, ContractInstanceWithAddress> = new Map();
+
+  public firstContractInstance: ContractInstanceWithAddress = SerializableContractInstance.default().withAddress(
+    AztecAddress.fromNumber(0),
+  );
+  public instanceSameClassAsFirstContract: ContractInstanceWithAddress =
+    SerializableContractInstance.default().withAddress(AztecAddress.fromNumber(0));
   public otherContractInstance: ContractInstanceWithAddress;
 
-  constructor() {
+  constructor(private skipContractDeployments: boolean = false) {
     this.bytecode = getAvmTestContractBytecode(this.fnName);
     this.fnSelector = getAvmTestContractFunctionSelector(this.fnName);
     this.publicFn = { bytecode: this.bytecode, selector: this.fnSelector };
-    this.contractClass = makeContractClassPublic(0, this.publicFn);
-    this.contractInstance = makeContractInstanceFromClassId(this.contractClass.id);
     this.bytecodeCommitment = computePublicBytecodeCommitment(this.bytecode);
+
+    // create enough unique classes to hit the limit (plus two extra)
+    for (let i = 0; i < MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS + 1; i++) {
+      const contractClass = makeContractClassPublic(/*seed=*/ i, this.publicFn);
+      const contractInstance = makeContractInstanceFromClassId(contractClass.id, /*seed=*/ i);
+      this.contractClasses.set(contractClass.id.toString(), contractClass);
+      this.contractInstances.set(contractInstance.address.toString(), contractInstance);
+      if (i === 0) {
+        this.firstContractInstance = contractInstance;
+      }
+    }
+    // a contract with the same class but different instance/address as the first contract
+    this.instanceSameClassAsFirstContract = makeContractInstanceFromClassId(
+      this.firstContractInstance.contractClassId,
+      /*seed=*/ 1000,
+    );
+
     // The values here should match those in `avm_simulator.test.ts`
     // Used for GETCONTRACTINSTANCE test
     this.otherContractInstance = new SerializableContractInstance({
@@ -178,6 +239,46 @@ export class MockedAvmTestContractDataSource {
     }).withAddress(AztecAddress.fromNumber(0x4444));
   }
 
+  async deployContracts(merkleTrees: MerkleTreeWriteOperations) {
+    if (!this.skipContractDeployments) {
+      for (const contractInstance of this.contractInstances.values()) {
+        const contractAddressNullifier = siloNullifier(
+          AztecAddress.fromNumber(DEPLOYER_CONTRACT_ADDRESS),
+          contractInstance.address.toField(),
+        );
+        await merkleTrees.batchInsert(MerkleTreeId.NULLIFIER_TREE, [contractAddressNullifier.toBuffer()], 0);
+      }
+
+      const instanceSameClassAsFirstContractNullifier = siloNullifier(
+        AztecAddress.fromNumber(DEPLOYER_CONTRACT_ADDRESS),
+        this.instanceSameClassAsFirstContract.address.toField(),
+      );
+      await merkleTrees.batchInsert(
+        MerkleTreeId.NULLIFIER_TREE,
+        [instanceSameClassAsFirstContractNullifier.toBuffer()],
+        0,
+      );
+
+      // other contract address used by the bulk test's GETCONTRACTINSTANCE test
+      const otherContractAddressNullifier = siloNullifier(
+        AztecAddress.fromNumber(DEPLOYER_CONTRACT_ADDRESS),
+        this.otherContractInstance.address.toField(),
+      );
+      await merkleTrees.batchInsert(MerkleTreeId.NULLIFIER_TREE, [otherContractAddressNullifier.toBuffer()], 0);
+    }
+  }
+
+  public static async create(
+    merkleTrees: MerkleTreeWriteOperations,
+    skipContractDeployments: boolean = false,
+  ): Promise<MockedAvmTestContractDataSource> {
+    const dataSource = new MockedAvmTestContractDataSource(skipContractDeployments);
+    if (!skipContractDeployments) {
+      await dataSource.deployContracts(merkleTrees);
+    }
+    return dataSource;
+  }
+
   getPublicFunction(_address: AztecAddress, _selector: FunctionSelector): Promise<PublicFunction> {
     return Promise.resolve(this.publicFn);
   }
@@ -186,8 +287,8 @@ export class MockedAvmTestContractDataSource {
     throw new Error('Method not implemented.');
   }
 
-  getContractClass(_id: Fr): Promise<ContractClassPublic> {
-    return Promise.resolve(this.contractClass);
+  getContractClass(id: Fr): Promise<ContractClassPublic | undefined> {
+    return Promise.resolve(this.contractClasses.get(id.toString()));
   }
 
   getBytecodeCommitment(_id: Fr): Promise<Fr> {
@@ -198,12 +299,17 @@ export class MockedAvmTestContractDataSource {
     return Promise.resolve();
   }
 
-  getContract(address: AztecAddress): Promise<ContractInstanceWithAddress> {
-    if (address.equals(this.contractInstance.address)) {
-      return Promise.resolve(this.contractInstance);
-    } else {
-      return Promise.resolve(this.otherContractInstance);
+  getContract(address: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
+    if (!this.skipContractDeployments) {
+      if (address.equals(this.otherContractInstance.address)) {
+        return Promise.resolve(this.otherContractInstance);
+      } else if (address.equals(this.instanceSameClassAsFirstContract.address)) {
+        return Promise.resolve(this.instanceSameClassAsFirstContract);
+      } else {
+        return Promise.resolve(this.contractInstances.get(address.toString()));
+      }
     }
+    return Promise.resolve(undefined);
   }
 
   getContractClassIds(): Promise<Fr[]> {
@@ -218,7 +324,7 @@ export class MockedAvmTestContractDataSource {
     return Promise.resolve(this.fnName);
   }
 
-  addContractArtifact(_address: AztecAddress, _contract: ContractArtifact): Promise<void> {
+  registerContractFunctionNames(_address: AztecAddress, _names: Record<string, string>): Promise<void> {
     return Promise.resolve();
   }
 }

package/src/public/public_db_sources.ts

@@ -22,7 +22,8 @@ import {
 import { computeL1ToL2MessageNullifier, computePublicDataTreeLeafSlot } from '@aztec/circuits.js/hash';
 import { createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
-import { ContractClassRegisteredEvent, ContractInstanceDeployedEvent } from '@aztec/protocol-contracts';
+import { ContractClassRegisteredEvent } from '@aztec/protocol-contracts/class-registerer';
+import { ContractInstanceDeployedEvent } from '@aztec/protocol-contracts/instance-deployer';
 import {
   type CommitmentsDB,
   MessageLoadOracleInputs,
@@ -207,7 +208,7 @@ export class WorldStateDB extends ContractsDataSourcePublicDB implements PublicS
     nullifier: Fr,
   ): Promise<NullifierMembershipWitness | undefined> {
     const timer = new Timer();
-    const index = await this.db.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBuffer());
+    const index = (await this.db.findLeafIndices(MerkleTreeId.NULLIFIER_TREE, [nullifier.toBuffer()]))[0];
     if (!index) {
       return undefined;
     }
@@ -240,7 +241,7 @@ export class WorldStateDB extends ContractsDataSourcePublicDB implements PublicS
   ): Promise<MessageLoadOracleInputs<typeof L1_TO_L2_MSG_TREE_HEIGHT>> {
     const timer = new Timer();
 
-    const messageIndex = await this.db.findLeafIndex(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, messageHash);
+    const messageIndex = (await this.db.findLeafIndices(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, [messageHash]))[0];
     if (messageIndex === undefined) {
       throw new Error(`No L1 to L2 message found for message hash ${messageHash.toString()}`);
     }
@@ -279,7 +280,7 @@ export class WorldStateDB extends ContractsDataSourcePublicDB implements PublicS
 
   public async getCommitmentIndex(commitment: Fr): Promise<bigint | undefined> {
     const timer = new Timer();
-    const index = await this.db.findLeafIndex(MerkleTreeId.NOTE_HASH_TREE, commitment);
+    const index = (await this.db.findLeafIndices(MerkleTreeId.NOTE_HASH_TREE, [commitment]))[0];
     this.logger.debug(`[DB] Fetched commitment index`, {
       eventName: 'public-db-access',
       duration: timer.ms(),
@@ -301,7 +302,7 @@ export class WorldStateDB extends ContractsDataSourcePublicDB implements PublicS
 
   public async getNullifierIndex(nullifier: Fr): Promise<bigint | undefined> {
     const timer = new Timer();
-    const index = await this.db.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBuffer());
+    const index = (await this.db.findLeafIndices(MerkleTreeId.NULLIFIER_TREE, [nullifier.toBuffer()]))[0];
     this.logger.debug(`[DB] Fetched nullifier index`, {
       eventName: 'public-db-access',
       duration: timer.ms(),