@aztec/simulator 0.0.1-commit.03f7ef2 → 0.0.1-commit.04d373f

package/src/public/fixtures/public_tx_simulation_tester.ts

@@ -1,11 +1,11 @@
-import { DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT } from '@aztec/constants';
+import { PUBLIC_TX_L2_GAS_OVERHEAD, TX_DA_GAS_OVERHEAD } from '@aztec/constants';
 import { asyncMap } from '@aztec/foundation/async-map';
 import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { type ContractArtifact, encodeArguments } from '@aztec/stdlib/abi';
 import { PublicSimulatorConfig, type PublicTxResult } from '@aztec/stdlib/avm';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { Gas, GasFees } from '@aztec/stdlib/gas';
+import { FALLBACK_TEARDOWN_DA_GAS_LIMIT, FALLBACK_TEARDOWN_L2_GAS_LIMIT, Gas, GasFees } from '@aztec/stdlib/gas';
 import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
 import { PublicCallRequest } from '@aztec/stdlib/kernel';
 import { GlobalVariables, PublicCallRequestWithCalldata, type Tx } from '@aztec/stdlib/tx';
@@ -112,6 +112,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     feePayer: AztecAddress = sender,
     /* need some unique first nullifier for note-nonce computations */
     privateInsertions: TestPrivateInsertions = { nonRevertible: { nullifiers: [new Fr(420000 + this.txCount)] } },
+    gasLimits?: Gas,
   ): Promise<Tx> {
     const setupCallRequests = await asyncMap(setupCalls, call =>
       this.#createPubicCallRequestForCall(call, call.sender ?? sender),
@@ -131,9 +132,13 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       teardownCallRequest,
       feePayer,
       /*gasUsedByPrivate*/ teardownCall
-        ? new Gas(DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT)
-        : Gas.empty(),
+        ? new Gas(
+            FALLBACK_TEARDOWN_DA_GAS_LIMIT + TX_DA_GAS_OVERHEAD,
+            FALLBACK_TEARDOWN_L2_GAS_LIMIT + PUBLIC_TX_L2_GAS_OVERHEAD,
+          )
+        : new Gas(TX_DA_GAS_OVERHEAD, PUBLIC_TX_L2_GAS_OVERHEAD),
       defaultGlobals(),
+      gasLimits,
     );
   }
 
@@ -146,8 +151,9 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     /* need some unique first nullifier for note-nonce computations */
     privateInsertions?: TestPrivateInsertions,
     txLabel: string = 'unlabeledTx',
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
-    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions);
+    const tx = await this.createTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions, gasLimits);
 
     await this.setFeePayerBalance(feePayer);
 
@@ -161,6 +167,8 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     }
     const avmResult = await this.simulator.simulate(tx, fullTxLabel);
 
+    await this.#recordBytecodeSizes(fullTxLabel, [...setupCalls, ...appCalls, ...(teardownCall ? [teardownCall] : [])]);
+
     // Something like this is often useful for debugging:
     //if (avmResult.revertReason) {
     //  // resolve / enrich revert reason
@@ -190,8 +198,18 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     teardownCall?: TestEnqueuedCall,
     feePayer?: AztecAddress,
     privateInsertions?: TestPrivateInsertions,
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
-    return await this.simulateTx(sender, setupCalls, appCalls, teardownCall, feePayer, privateInsertions, txLabel);
+    return await this.simulateTx(
+      sender,
+      setupCalls,
+      appCalls,
+      teardownCall,
+      feePayer,
+      privateInsertions,
+      txLabel,
+      gasLimits,
+    );
   }
 
   /**
@@ -209,6 +227,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     teardownCall?: TestEnqueuedCall,
     feePayer?: AztecAddress,
     privateInsertions?: TestPrivateInsertions,
+    gasLimits?: Gas,
   ): Promise<PublicTxResult> {
     return await this.simulateTxWithLabel(
       txLabel,
@@ -218,6 +237,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       teardownCall,
       feePayer,
       privateInsertions,
+      gasLimits,
     );
   }
 
@@ -277,6 +297,27 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
 
     return new PublicCallRequestWithCalldata(request, calldata);
   }
+
+  // WARNING: Deduplicates by artifact name, so two different artifacts with the same name
+  // in a single tx would only record the first one's bytecode size.
+  async #recordBytecodeSizes(txLabel: string, calls: TestEnqueuedCall[]) {
+    const seenArtifactNames = new Set<string>();
+    for (const call of calls) {
+      const artifact = await this.contractDataSource.getContractArtifact(call.address);
+      if (!artifact || seenArtifactNames.has(artifact.name)) {
+        continue;
+      }
+      seenArtifactNames.add(artifact.name);
+      const instance = await this.contractDataSource.getContract(call.address);
+      if (!instance) {
+        continue;
+      }
+      const contractClass = await this.contractDataSource.getContractClass(instance.currentContractClassId);
+      if (contractClass) {
+        this.metrics.recordBytecodeSize(txLabel, artifact.name, contractClass.packedBytecode.length);
+      }
+    }
+  }
 }
 
 export function defaultGlobals() {

package/src/public/fixtures/utils.ts

@@ -2,10 +2,8 @@ import {
   CONTRACT_CLASS_PUBLISHED_MAGIC_VALUE,
   CONTRACT_CLASS_REGISTRY_CONTRACT_ADDRESS,
   CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
-  DEFAULT_DA_GAS_LIMIT,
-  DEFAULT_L2_GAS_LIMIT,
-  DEFAULT_TEARDOWN_DA_GAS_LIMIT,
-  DEFAULT_TEARDOWN_L2_GAS_LIMIT,
+  MAX_PROCESSABLE_DA_GAS_PER_CHECKPOINT,
+  MAX_PROCESSABLE_L2_GAS,
   PRIVATE_LOG_SIZE_IN_FIELDS,
 } from '@aztec/constants';
 import { padArrayEnd } from '@aztec/foundation/collection';
@@ -14,7 +12,13 @@ import { CONTRACT_INSTANCE_PUBLISHED_EVENT_TAG } from '@aztec/protocol-contracts
 import { bufferAsFields } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import type { ContractClassPublic, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
-import { Gas, GasFees, GasSettings } from '@aztec/stdlib/gas';
+import {
+  FALLBACK_TEARDOWN_DA_GAS_LIMIT,
+  FALLBACK_TEARDOWN_L2_GAS_LIMIT,
+  Gas,
+  GasFees,
+  GasSettings,
+} from '@aztec/stdlib/gas';
 import { siloNullifier } from '@aztec/stdlib/hash';
 import {
   LogHash,
@@ -62,13 +66,14 @@ export async function createTxForPublicCalls(
   feePayer = AztecAddress.zero(),
   gasUsedByPrivate: Gas = Gas.empty(),
   globals: GlobalVariables = GlobalVariables.empty(),
+  gasLimits?: Gas,
 ): Promise<Tx> {
   assert(
     setupCallRequests.length > 0 || appCallRequests.length > 0 || teardownCallRequest !== undefined,
     "Can't create public tx with no enqueued calls",
   );
   // use max limits
-  const gasLimits = new Gas(DEFAULT_DA_GAS_LIMIT, DEFAULT_L2_GAS_LIMIT);
+  gasLimits = gasLimits ?? new Gas(MAX_PROCESSABLE_DA_GAS_PER_CHECKPOINT, MAX_PROCESSABLE_L2_GAS);
 
   const forPublic = PartialPrivateTailPublicInputsForPublic.empty();
 
@@ -128,19 +133,19 @@ export async function createTxForPublicCalls(
 
   const maxFeesPerGas = feePayer.isZero() ? GasFees.empty() : new GasFees(10, 10);
   const teardownGasLimits = teardownCallRequest
-    ? new Gas(DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT)
+    ? new Gas(FALLBACK_TEARDOWN_DA_GAS_LIMIT, FALLBACK_TEARDOWN_L2_GAS_LIMIT)
     : Gas.empty();
   const gasSettings = new GasSettings(gasLimits, teardownGasLimits, maxFeesPerGas, GasFees.empty());
   const txContext = new TxContext(Fr.zero(), Fr.zero(), gasSettings);
   const header = BlockHeader.empty({ globalVariables: globals });
   const constantData = new TxConstantData(header, txContext, Fr.zero(), Fr.zero());
-  const includeByTimestamp = 0n; // Not used in the simulator.
+  const expirationTimestamp = 0n; // Not used in the simulator.
 
   const txData = new PrivateKernelTailCircuitPublicInputs(
     constantData,
     /*gasUsed=*/ gasUsedByPrivate,
     feePayer,
-    includeByTimestamp,
+    expirationTimestamp,
     forPublic,
   );
 
@@ -163,7 +168,7 @@ export async function createTxForPrivateOnly(
   gasUsedByPrivate: Gas = new Gas(10, 10),
 ): Promise<Tx> {
   // use max limits
-  const gasLimits = new Gas(DEFAULT_DA_GAS_LIMIT, DEFAULT_L2_GAS_LIMIT);
+  const gasLimits = new Gas(MAX_PROCESSABLE_DA_GAS_PER_CHECKPOINT, MAX_PROCESSABLE_L2_GAS);
 
   const forRollup = PartialPrivateTailPublicInputsForRollup.empty();
 
@@ -171,13 +176,13 @@ export async function createTxForPrivateOnly(
   const gasSettings = new GasSettings(gasLimits, Gas.empty(), maxFeesPerGas, GasFees.empty());
   const txContext = new TxContext(Fr.zero(), Fr.zero(), gasSettings);
   const constantData = new TxConstantData(BlockHeader.empty(), txContext, Fr.zero(), Fr.zero());
-  const includeByTimestamp = 0n; // Not used in the simulator.
+  const expirationTimestamp = 0n; // Not used in the simulator.
 
   const txData = new PrivateKernelTailCircuitPublicInputs(
     constantData,
     /*gasUsed=*/ gasUsedByPrivate,
     feePayer,
-    includeByTimestamp,
+    expirationTimestamp,
     /*forPublic=*/ undefined,
     forRollup,
   );
@@ -228,17 +233,15 @@ export async function addNewContractInstanceToTx(
   contractInstance: ContractInstanceWithAddress,
   skipNullifierInsertion = false,
 ) {
-  // can't use publicKeys.toFields() because it includes isInfinite which
-  // is not broadcast in such private logs
+  // Only ivpk_m is broadcast as a point (x, y); the other five keys are hashes.
   const publicKeysAsFields = [
-    contractInstance.publicKeys.masterNullifierPublicKey.x,
-    contractInstance.publicKeys.masterNullifierPublicKey.y,
-    contractInstance.publicKeys.masterIncomingViewingPublicKey.x,
-    contractInstance.publicKeys.masterIncomingViewingPublicKey.y,
-    contractInstance.publicKeys.masterOutgoingViewingPublicKey.x,
-    contractInstance.publicKeys.masterOutgoingViewingPublicKey.y,
-    contractInstance.publicKeys.masterTaggingPublicKey.x,
-    contractInstance.publicKeys.masterTaggingPublicKey.y,
+    contractInstance.publicKeys.npkMHash,
+    contractInstance.publicKeys.ivpkM.x,
+    contractInstance.publicKeys.ivpkM.y,
+    contractInstance.publicKeys.ovpkMHash,
+    contractInstance.publicKeys.tpkMHash,
+    contractInstance.publicKeys.mspkMHash,
+    contractInstance.publicKeys.fbpkMHash,
   ];
   const logFields = [
     CONTRACT_INSTANCE_PUBLISHED_EVENT_TAG,
@@ -247,6 +250,7 @@ export async function addNewContractInstanceToTx(
     new Fr(contractInstance.salt),
     contractInstance.currentContractClassId,
     contractInstance.initializationHash,
+    contractInstance.immutablesHash,
     ...publicKeysAsFields,
     contractInstance.deployer.toField(),
   ];

package/src/public/fuzzing/avm_fuzzer_simulator.ts

@@ -4,6 +4,7 @@ import {
   MAX_NOTE_HASHES_PER_TX,
   MAX_NULLIFIERS_PER_TX,
   MAX_PRIVATE_LOGS_PER_TX,
+  MAX_PROTOCOL_CONTRACTS,
 } from '@aztec/constants';
 import { padArrayEnd } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/curves/bn254';
@@ -19,8 +20,17 @@ import {
 import { PrivateLog } from '@aztec/stdlib/logs';
 import { ScopedL2ToL1Message } from '@aztec/stdlib/messaging';
 import { ChonkProof } from '@aztec/stdlib/proofs';
-import type { MerkleTreeWriteOperations } from '@aztec/stdlib/trees';
-import { BlockHeader, GlobalVariables, HashedValues, Tx, TxConstantData, TxContext, TxHash } from '@aztec/stdlib/tx';
+import { MerkleTreeId, type MerkleTreeWriteOperations, PublicDataTreeLeaf } from '@aztec/stdlib/trees';
+import {
+  BlockHeader,
+  GlobalVariables,
+  HashedValues,
+  ProtocolContracts,
+  Tx,
+  TxConstantData,
+  TxContext,
+  TxHash,
+} from '@aztec/stdlib/tx';
 import type { NativeWorldStateService } from '@aztec/world-state';
 
 import { BaseAvmSimulationTester } from '../avm/fixtures/base_avm_simulation_tester.js';
@@ -40,6 +50,9 @@ export class FuzzerSimulationRequest {
     public readonly globals: GlobalVariables,
     public readonly contractClasses: any[], // Raw, processed by addContractClassFromCpp
     public readonly contractInstances: [any, any][], // Raw pairs [address, instance]
+    public readonly publicDataWrites: any[], // Raw public data tree writes to apply before simulation
+    public readonly noteHashes: any[], // Raw note hashes to apply before simulation
+    public readonly protocolContracts: ProtocolContracts, // Protocol contracts mapping from C++
   ) {}
 
   static fromPlainObject(obj: any): FuzzerSimulationRequest {
@@ -53,6 +66,9 @@ export class FuzzerSimulationRequest {
       GlobalVariables.fromPlainObject(obj.globals),
       obj.contractClasses,
       obj.contractInstances,
+      obj.publicDataWrites ?? [],
+      obj.noteHashes ?? [],
+      ProtocolContracts.fromPlainObject(obj.protocolContracts),
     );
   }
 }
@@ -130,7 +146,7 @@ async function createTxFromHint(cppTx: AvmTxHint): Promise<Tx> {
     constants,
     cppTx.gasUsedByPrivate,
     cppTx.feePayer,
-    0n, // includeByTimestamp
+    0n, // expirationTimestamp
     forPublic,
     undefined, // forRollup - not needed for public simulation
   );
@@ -181,16 +197,23 @@ export class AvmFuzzerSimulator extends BaseAvmSimulationTester {
     merkleTrees: MerkleTreeWriteOperations,
     contractDataSource: SimpleContractDataSource,
     globals: GlobalVariables,
+    protocolContracts: ProtocolContracts,
   ) {
     super(contractDataSource, merkleTrees);
     const contractsDb = new PublicContractsDB(contractDataSource);
-    this.simulator = new PublicTxSimulator(merkleTrees, contractsDb, globals, {
-      skipFeeEnforcement: false,
-      collectDebugLogs: false,
-      collectHints: false,
-      collectStatistics: false,
-      collectCallMetadata: false,
-    });
+    this.simulator = new PublicTxSimulator(
+      merkleTrees,
+      contractsDb,
+      globals,
+      {
+        skipFeeEnforcement: false,
+        collectDebugLogs: false,
+        collectHints: false,
+        collectStatistics: false,
+        collectCallMetadata: false,
+      },
+      protocolContracts,
+    );
   }
 
   /**
@@ -199,10 +222,11 @@ export class AvmFuzzerSimulator extends BaseAvmSimulationTester {
   public static async create(
     worldStateService: NativeWorldStateService,
     globals: GlobalVariables,
+    protocolContracts: ProtocolContracts,
   ): Promise<AvmFuzzerSimulator> {
     const contractDataSource = new SimpleContractDataSource();
     const merkleTrees = await worldStateService.fork();
-    return new AvmFuzzerSimulator(merkleTrees, contractDataSource, globals);
+    return new AvmFuzzerSimulator(merkleTrees, contractDataSource, globals, protocolContracts);
   }
 
   /**
@@ -230,11 +254,35 @@ export class AvmFuzzerSimulator extends BaseAvmSimulationTester {
 
   /**
    * Add a contract instance from C++ raw msgpack data.
-   * This also inserts the contract address nullifier into the nullifier tree.
+   * This also inserts the contract address nullifier into the nullifier tree,
+   * unless the address is a protocol canonical address (1-11).
    */
   public async addContractInstanceFromCpp(rawAddress: any, rawInstance: any): Promise<void> {
     const address = AztecAddress.fromPlainObject(rawAddress);
     const instance = contractInstanceWithAddressFromPlainObject(address, rawInstance);
-    await this.addContractInstance(instance);
+    // Protocol canonical addresses (1-11) should not have nullifiers inserted
+    const isProtocolCanonicalAddress = address.toBigInt() <= MAX_PROTOCOL_CONTRACTS && address.toBigInt() >= 1n;
+    await this.addContractInstance(instance, /* skipNullifierInsertion */ isProtocolCanonicalAddress);
+  }
+
+  /**
+   * Apply public data tree writes from C++ raw msgpack data.
+   * This is used to pre-populate the public data tree before simulation (e.g., for bytecode upgrades).
+   */
+  public async applyPublicDataWrites(rawWrites: any[]): Promise<void> {
+    for (const rawWrite of rawWrites) {
+      const leaf = PublicDataTreeLeaf.fromPlainObject(rawWrite);
+      await this.merkleTrees.sequentialInsert(MerkleTreeId.PUBLIC_DATA_TREE, [leaf.toBuffer()]);
+    }
+  }
+
+  /**
+   * Apply note hashes from C++ raw msgpack data.
+   * This is used to pre-populate the note hash tree before simulation.
+   */
+  public async applyNoteHashes(rawNoteHashes: any[]): Promise<void> {
+    const paddingLeaves = MAX_NOTE_HASHES_PER_TX - (rawNoteHashes.length % MAX_NOTE_HASHES_PER_TX);
+    const paddedNoteHashes = [...rawNoteHashes, ...Array(paddingLeaves).fill(Fr.ZERO)];
+    await this.merkleTrees.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, paddedNoteHashes);
   }
 }

package/src/public/fuzzing/avm_simulator_bin.ts

@@ -3,17 +3,30 @@ import { EthAddress } from '@aztec/foundation/eth-address';
 import {
   AvmCircuitPublicInputs,
   type AvmTxHint,
+  PublicTxEffect,
   deserializeFromMessagePack,
   serializeWithMessagePack,
 } from '@aztec/stdlib/avm';
-import { GlobalVariables, TreeSnapshots } from '@aztec/stdlib/tx';
+import { GlobalVariables, ProtocolContracts, TreeSnapshots } from '@aztec/stdlib/tx';
 import { NativeWorldStateService } from '@aztec/world-state';
 
-import { writeSync } from 'fs';
 import { createInterface } from 'readline';
 
 import { AvmFuzzerSimulator, FuzzerSimulationRequest } from './avm_fuzzer_simulator.js';
 
+/** Write data to stdout, letting Node handle buffering. */
+function writeOutput(data: string): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    process.stdout.write(data, err => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve();
+      }
+    });
+  });
+}
+
 // This cache holds opened world states to avoid reopening them for each invocation.
 // It's a map so that in the future we could support multiple world states (if we had multiple fuzzers).
 const worldStateCache = new Map<string, NativeWorldStateService>();
@@ -43,12 +56,23 @@ async function simulateWithFuzzer(
   globals: GlobalVariables,
   rawContractClasses: any[], // Replace these when we are moving contract classes to TS
   rawContractInstances: [any, any][], // Replace these when we are moving contract instances to TS
-): Promise<{ reverted: boolean; output: Fr[]; revertReason?: string; publicInputs: AvmCircuitPublicInputs }> {
+  rawPublicDataWrites: any[], // Public data tree writes to apply before simulation
+  rawNoteHashes: any[], // Note hashes to apply before simulation
+  protocolContracts: ProtocolContracts, // Protocol contracts mapping from C++
+): Promise<{
+  reverted: boolean;
+  output: Fr[];
+  revertReason?: string;
+  publicInputs: AvmCircuitPublicInputs;
+  publicTxEffect: PublicTxEffect;
+}> {
   const worldStateService = await openExistingWorldState(dataDir, mapSizeKb);
 
-  const simulator = await AvmFuzzerSimulator.create(worldStateService, globals);
+  const simulator = await AvmFuzzerSimulator.create(worldStateService, globals, protocolContracts);
+
+  await simulator.applyNoteHashes(rawNoteHashes);
 
-  // Register contract classes from C++
+  // Register contract classes from C++ (must happen before public data writes to match C++ order)
   for (const rawClass of rawContractClasses) {
     await simulator.addContractClassFromCpp(rawClass);
   }
@@ -58,6 +82,10 @@ async function simulateWithFuzzer(
     await simulator.addContractInstanceFromCpp(rawAddress, rawInstance);
   }
 
+  // Apply public data writes after contract registration (e.g., for bytecode upgrades)
+  // This must happen last to match C++ setup_fuzzer_state ordering
+  await simulator.applyPublicDataWrites(rawPublicDataWrites);
+
   const result = await simulator.simulate(txHint);
 
   const output = result
@@ -69,6 +97,7 @@ async function simulateWithFuzzer(
     output,
     revertReason: result.findRevertReason()?.message,
     publicInputs: result.publicInputs!,
+    publicTxEffect: result.publicTxEffect,
   };
 }
 
@@ -87,6 +116,9 @@ async function execute(base64Line: string): Promise<void> {
       request.globals,
       request.contractClasses,
       request.contractInstances,
+      request.publicDataWrites,
+      request.noteHashes,
+      request.protocolContracts,
     );
 
     // Serialize the result to msgpack and encode it in base64 for output
@@ -95,17 +127,20 @@ async function execute(base64Line: string): Promise<void> {
       output: result.output,
       revertReason: result.revertReason ?? '',
       endTreeSnapshots: result.publicInputs.endTreeSnapshots,
+      publicTxEffect: result.publicTxEffect,
     });
-    writeSync(process.stdout.fd, resultBuffer.toString('base64') + '\n');
+    const base64Response = resultBuffer.toString('base64') + '\n';
+    await writeOutput(base64Response);
   } catch (error: any) {
     // If we error, treat as reverted
     const errorResult = serializeWithMessagePack({
       reverted: true,
-      output: [] as string[],
+      output: [] as Fr[],
       revertReason: `Unexpected Error ${error.message}`,
       endTreeSnapshots: TreeSnapshots.empty(),
+      publicTxEffect: PublicTxEffect.empty(),
     });
-    writeSync(process.stdout.fd, errorResult.toString('base64') + '\n');
+    await writeOutput(errorResult.toString('base64') + '\n');
   }
 }
 

package/src/public/hinting_db_sources.ts

@@ -83,6 +83,7 @@ export class HintingPublicContractsDB implements PublicContractsDBInterface {
           instance.currentContractClassId,
           instance.originalContractClassId,
           instance.initializationHash,
+          instance.immutablesHash,
           instance.publicKeys,
         ),
       );
@@ -410,12 +411,12 @@ export class HintingMerkleWriteOperations implements MerkleTreeWriteOperations {
     }
   }
 
-  public async createCheckpoint(): Promise<void> {
+  public async createCheckpoint(): Promise<number> {
     const actionCounter = this.checkpointActionCounter++;
     const oldCheckpointId = this.getCurrentCheckpointId();
     const treesStateHash = await this.getTreesStateHash();
 
-    await this.db.createCheckpoint();
+    const depth = await this.db.createCheckpoint();
     this.checkpointStack.push(this.nextCheckpointId++);
     const newCheckpointId = this.getCurrentCheckpointId();
 
@@ -424,14 +425,16 @@ export class HintingMerkleWriteOperations implements MerkleTreeWriteOperations {
     HintingMerkleWriteOperations.log.trace(
       `[createCheckpoint:${actionCounter}] Checkpoint evolved ${oldCheckpointId} -> ${newCheckpointId} at trees state ${treesStateHash}.`,
     );
+
+    return depth;
   }
 
-  public commitAllCheckpoints(): Promise<void> {
-    throw new Error('commitAllCheckpoints is not supported in HintingMerkleWriteOperations.');
+  public commitAllCheckpointsTo(_depth: number): Promise<void> {
+    throw new Error('commitAllCheckpointsTo is not supported in HintingMerkleWriteOperations.');
   }
 
-  public revertAllCheckpoints(): Promise<void> {
-    throw new Error('revertAllCheckpoints is not supported in HintingMerkleWriteOperations.');
+  public revertAllCheckpointsTo(_depth: number): Promise<void> {
+    throw new Error('revertAllCheckpointsTo is not supported in HintingMerkleWriteOperations.');
   }
 
   public async commitCheckpoint(): Promise<void> {
@@ -572,7 +575,7 @@ export class HintingMerkleWriteOperations implements MerkleTreeWriteOperations {
     return await this.db.close();
   }
 
-  async [Symbol.dispose](): Promise<void> {
+  async [Symbol.asyncDispose](): Promise<void> {
     await this.close();
   }
 

package/src/public/public_db_sources.ts

@@ -5,7 +5,7 @@ import {
   PUBLIC_DATA_SUBTREE_HEIGHT,
 } from '@aztec/constants';
 import { Fr } from '@aztec/foundation/curves/bn254';
-import { createLogger } from '@aztec/foundation/log';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import { ContractClassPublishedEvent } from '@aztec/protocol-contracts/class-registry';
 import { ContractInstancePublishedEvent } from '@aztec/protocol-contracts/instance-registry';
@@ -46,14 +46,20 @@ import { L1ToL2MessageIndexOutOfRangeError, NoteHashIndexOutOfRangeError } from
 export class PublicContractsDB implements PublicContractsDBInterface {
   private contractStateStack: ContractsDbCheckpoint[] = [new ContractsDbCheckpoint()];
 
-  private log = createLogger('simulator:contracts-data-source');
+  private log: Logger;
 
-  constructor(private dataSource: ContractDataSource) {}
+  constructor(
+    private dataSource: ContractDataSource,
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('simulator:contracts-data-source', bindings);
+  }
 
-  public async addContracts(contractDeploymentData: ContractDeploymentData): Promise<void> {
+  /** Parses raw log data from the C++/NAPI bridge and inserts the resulting contracts into the current checkpoint. */
+  public addContractsFromLogs(contractDeploymentData: ContractDeploymentData): void {
     const currentState = this.getCurrentState();
 
-    await this.addContractClassesFromEvents(
+    this.addContractClassesFromEvents(
       ContractClassPublishedEvent.extractContractClassEvents(contractDeploymentData.getContractClassLogs()),
       currentState,
     );
@@ -64,10 +70,18 @@ export class PublicContractsDB implements PublicContractsDBInterface {
     );
   }
 
-  public async addNewContracts(tx: Tx): Promise<void> {
+  public addNewContracts(tx: Tx): void {
     const contractDeploymentData = AllContractDeploymentData.fromTx(tx);
-    await this.addContracts(contractDeploymentData.getNonRevertibleContractDeploymentData());
-    await this.addContracts(contractDeploymentData.getRevertibleContractDeploymentData());
+    this.addContractsFromLogs(contractDeploymentData.getNonRevertibleContractDeploymentData());
+    this.addContractsFromLogs(contractDeploymentData.getRevertibleContractDeploymentData());
+  }
+
+  /** Inserts typed contract instances directly into the current checkpoint. */
+  public addContracts(contractInstances?: ContractInstanceWithAddress[]): void {
+    const currentState = this.getCurrentState();
+    for (const instance of contractInstances ?? []) {
+      currentState.addInstance(instance.address, instance);
+    }
   }
 
   /**
@@ -76,7 +90,7 @@ export class PublicContractsDB implements PublicContractsDBInterface {
    */
   public createCheckpoint(): void {
     const currentState = this.getCurrentState();
-    const newState = currentState.deepCopy();
+    const newState = currentState.fork();
     this.contractStateStack.push(newState);
   }
 
@@ -169,17 +183,15 @@ export class PublicContractsDB implements PublicContractsDBInterface {
     return await this.dataSource.getDebugFunctionName(address, selector);
   }
 
-  private async addContractClassesFromEvents(
+  private addContractClassesFromEvents(
     contractClassEvents: ContractClassPublishedEvent[],
     state: ContractsDbCheckpoint,
   ) {
-    await Promise.all(
-      contractClassEvents.map(async (event: ContractClassPublishedEvent) => {
-        this.log.debug(`Adding class ${event.contractClassId.toString()} to contract state`);
-        const contractClass = await event.toContractClassPublic();
-        state.addClass(event.contractClassId, contractClass);
-      }),
-    );
+    for (const event of contractClassEvents) {
+      this.log.debug(`Adding class ${event.contractClassId.toString()} to contract state`);
+      const contractClass = event.toContractClassPublic();
+      state.addClass(event.contractClassId, contractClass);
+    }
   }
 
   private addContractInstancesFromEvents(
@@ -208,9 +220,14 @@ export class PublicContractsDB implements PublicContractsDBInterface {
  * to decide whether to use hints or not (same with tracing, etc).
  */
 export class PublicTreesDB implements PublicStateDBInterface {
-  private logger = createLogger('simulator:public-trees-db');
+  private logger: Logger;
 
-  constructor(private readonly db: MerkleTreeWriteOperations) {}
+  constructor(
+    private readonly db: MerkleTreeWriteOperations,
+    bindings?: LoggerBindings,
+  ) {
+    this.logger = createLogger('simulator:public-trees-db', bindings);
+  }
 
   public async storageRead(contract: AztecAddress, slot: Fr): Promise<Fr> {
     const timer = new Timer();