@aztec/simulator 0.0.1-commit.5daedc8 → 0.0.1-commit.6201a7b05

package/src/public/fuzzing/avm_fuzzer_simulator.ts

@@ -0,0 +1,288 @@
+import {
+  MAX_ENQUEUED_CALLS_PER_TX,
+  MAX_L2_TO_L1_MSGS_PER_TX,
+  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';
+import { AvmTxHint, type PublicTxResult } from '@aztec/stdlib/avm';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { contractClassPublicFromPlainObject, contractInstanceWithAddressFromPlainObject } from '@aztec/stdlib/contract';
+import {
+  PartialPrivateTailPublicInputsForPublic,
+  PrivateKernelTailCircuitPublicInputs,
+  PrivateToPublicAccumulatedData,
+  PublicCallRequest,
+} from '@aztec/stdlib/kernel';
+import { PrivateLog } from '@aztec/stdlib/logs';
+import { ScopedL2ToL1Message } from '@aztec/stdlib/messaging';
+import { ChonkProof } from '@aztec/stdlib/proofs';
+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';
+import { SimpleContractDataSource } from '../fixtures/simple_contract_data_source.js';
+import { PublicContractsDB } from '../public_db_sources.js';
+import { PublicTxSimulator } from '../public_tx_simulator/public_tx_simulator.js';
+
+/**
+ * Request structure for fuzzer simulation communication from C++.
+ * Matches the C++ FuzzerSimulationRequest struct
+ */
+export class FuzzerSimulationRequest {
+  constructor(
+    public readonly wsDataDir: string,
+    public readonly wsMapSizeKb: number,
+    public readonly tx: AvmTxHint,
+    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 {
+    if (obj instanceof FuzzerSimulationRequest) {
+      return obj;
+    }
+    return new FuzzerSimulationRequest(
+      obj.wsDataDir,
+      obj.wsMapSizeKb,
+      AvmTxHint.fromPlainObject(obj.tx),
+      GlobalVariables.fromPlainObject(obj.globals),
+      obj.contractClasses,
+      obj.contractInstances,
+      obj.publicDataWrites ?? [],
+      obj.noteHashes ?? [],
+      ProtocolContracts.fromPlainObject(obj.protocolContracts),
+    );
+  }
+}
+
+/**
+ * Creates a TypeScript Tx object from a deserialized C++ Tx (AvmTxHint-like structure).
+ * This allows using PublicTxSimulator.simulate() with fuzzer-generated transactions.
+ */
+async function createTxFromHint(cppTx: AvmTxHint): Promise<Tx> {
+  // Create TxHash from the C++ tx hash string
+  if (!cppTx.hash) {
+    throw new Error(`cppTx.hash is undefined. Keys: ${Object.keys(cppTx || {}).join(', ')}`);
+  }
+  const txHash = TxHash.fromString(cppTx.hash);
+
+  // Extract PublicCallRequest instances from enqueued calls
+  const setupCallRequests = cppTx.setupEnqueuedCalls.map(call => call.request);
+  const paddedSetupCalls = padArrayEnd(setupCallRequests, PublicCallRequest.empty(), MAX_ENQUEUED_CALLS_PER_TX);
+
+  const appLogicCallRequests = cppTx.appLogicEnqueuedCalls.map(call => call.request);
+  const paddedAppLogicCalls = padArrayEnd(appLogicCallRequests, PublicCallRequest.empty(), MAX_ENQUEUED_CALLS_PER_TX);
+
+  // Build non-revertible accumulated data from C++ tx
+  const emptyNonRevertible = PrivateToPublicAccumulatedData.empty();
+  const nonRevertibleAccumulatedData = new PrivateToPublicAccumulatedData(
+    padArrayEnd(cppTx.nonRevertibleAccumulatedData.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
+    padArrayEnd(cppTx.nonRevertibleAccumulatedData.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX),
+    padArrayEnd(
+      cppTx.nonRevertibleAccumulatedData.l2ToL1Messages,
+      ScopedL2ToL1Message.empty(),
+      MAX_L2_TO_L1_MSGS_PER_TX,
+    ),
+    padArrayEnd(cppTx.nonRevertibleContractDeploymentData.privateLogs, PrivateLog.empty(), MAX_PRIVATE_LOGS_PER_TX),
+    emptyNonRevertible.contractClassLogsHashes,
+    paddedSetupCalls,
+  );
+
+  // Build revertible accumulated data from C++ tx
+  const emptyRevertible = PrivateToPublicAccumulatedData.empty();
+  const revertibleAccumulatedData = new PrivateToPublicAccumulatedData(
+    padArrayEnd(cppTx.revertibleAccumulatedData.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
+    padArrayEnd(cppTx.revertibleAccumulatedData.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX),
+    padArrayEnd(cppTx.revertibleAccumulatedData.l2ToL1Messages, ScopedL2ToL1Message.empty(), MAX_L2_TO_L1_MSGS_PER_TX),
+    padArrayEnd(cppTx.revertibleContractDeploymentData.privateLogs, PrivateLog.empty(), MAX_PRIVATE_LOGS_PER_TX),
+    emptyRevertible.contractClassLogsHashes,
+    paddedAppLogicCalls,
+  );
+
+  // Build teardown call request (if exists)
+  const teardownCallRequest = cppTx.teardownEnqueuedCall?.request ?? PublicCallRequest.empty();
+
+  // Create forPublic structure
+  const forPublic = new PartialPrivateTailPublicInputsForPublic(
+    nonRevertibleAccumulatedData,
+    revertibleAccumulatedData,
+    teardownCallRequest,
+  );
+
+  // Build TxContext - gasSettings is already a proper GasSettings after AvmTxHint.fromPlainObject
+  const txContext = new TxContext(
+    Fr.ZERO, // chainId - this is fine because simulation actually reads from globalVariables not here
+    Fr.ZERO, // version - this is fine because simulation actually reads from globalVariables not here
+    cppTx.gasSettings,
+  );
+
+  // Build TxConstantData
+  const constants = new TxConstantData(
+    BlockHeader.empty(), // anchorBlockHeader (unused in simulation)
+    txContext,
+    Fr.ZERO, // vkTreeRoot - not needed for public simulation
+    Fr.ZERO, // protocolContractsHash - not needed for public simulation
+  );
+
+  const data = new PrivateKernelTailCircuitPublicInputs(
+    constants,
+    cppTx.gasUsedByPrivate,
+    cppTx.feePayer,
+    0n, // expirationTimestamp
+    forPublic,
+    undefined, // forRollup - not needed for public simulation
+  );
+
+  // Build publicFunctionCalldata from all enqueued calls
+  // Calldata is already Fr[] after AvmTxHint.fromPlainObject
+  const publicFunctionCalldata: HashedValues[] = [];
+
+  // Add setup calls
+  for (const call of cppTx.setupEnqueuedCalls || []) {
+    publicFunctionCalldata.push(await HashedValues.fromCalldata(call.calldata));
+  }
+
+  // Add app logic calls
+  for (const call of cppTx.appLogicEnqueuedCalls || []) {
+    publicFunctionCalldata.push(await HashedValues.fromCalldata(call.calldata));
+  }
+
+  // Add teardown call if present
+  if (cppTx.teardownEnqueuedCall) {
+    publicFunctionCalldata.push(await HashedValues.fromCalldata(cppTx.teardownEnqueuedCall.calldata));
+  }
+
+  // Extract contract class log fields from ContractDeploymentData
+  const contractClassLogFields = [
+    ...cppTx.nonRevertibleContractDeploymentData.contractClassLogs.map(log => log.fields),
+    ...cppTx.revertibleContractDeploymentData.contractClassLogs.map(log => log.fields),
+  ];
+
+  // Create the Tx
+  return new Tx(
+    txHash,
+    data,
+    ChonkProof.empty(), // No real proof needed for simulation
+    contractClassLogFields,
+    publicFunctionCalldata,
+  );
+}
+
+/**
+ * A simulator class for the AVM fuzzer that extends BaseAvmSimulationTester.
+ * It provides methods for registering contracts from C++ msgpack data and simulating transactions.
+ */
+export class AvmFuzzerSimulator extends BaseAvmSimulationTester {
+  private simulator: PublicTxSimulator;
+
+  constructor(
+    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,
+      },
+      protocolContracts,
+    );
+  }
+
+  /**
+   * Static factory method to create an AvmFuzzerSimulator.
+   */
+  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, protocolContracts);
+  }
+
+  /**
+   * Simulate a transaction from a C++ AvmTxHint.
+   */
+  public async simulate(txHint: AvmTxHint): Promise<PublicTxResult> {
+    // Compute fee from gas limits and max fees per gas (upper bound on fee)
+    const totalFee =
+      BigInt(txHint.gasSettings.gasLimits.daGas) * txHint.gasSettings.maxFeesPerGas.feePerDaGas +
+      BigInt(txHint.gasSettings.gasLimits.l2Gas) * txHint.gasSettings.maxFeesPerGas.feePerL2Gas;
+
+    await this.setFeePayerBalance(txHint.feePayer, new Fr(totalFee));
+
+    const tx = await createTxFromHint(txHint);
+    return await this.simulator.simulate(tx);
+  }
+
+  /**
+   * Add a contract class from C++ raw msgpack data.
+   */
+  public async addContractClassFromCpp(rawClass: any): Promise<void> {
+    const contractClass = contractClassPublicFromPlainObject(rawClass);
+    await this.contractDataSource.addContractClass(contractClass);
+  }
+
+  /**
+   * Add a contract instance from C++ raw msgpack data.
+   * 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);
+    // 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

@@ -0,0 +1,175 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { EthAddress } from '@aztec/foundation/eth-address';
+import {
+  AvmCircuitPublicInputs,
+  type AvmTxHint,
+  PublicTxEffect,
+  deserializeFromMessagePack,
+  serializeWithMessagePack,
+} from '@aztec/stdlib/avm';
+import { GlobalVariables, ProtocolContracts, TreeSnapshots } from '@aztec/stdlib/tx';
+import { NativeWorldStateService } from '@aztec/world-state';
+
+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>();
+
+async function openExistingWorldState(dataDir: string, mapSizeKb: number): Promise<NativeWorldStateService> {
+  const cached = worldStateCache.get(dataDir);
+  if (cached) {
+    return cached;
+  }
+
+  const ws = await NativeWorldStateService.new(EthAddress.ZERO, dataDir, {
+    archiveTreeMapSizeKb: mapSizeKb,
+    nullifierTreeMapSizeKb: mapSizeKb,
+    noteHashTreeMapSizeKb: mapSizeKb,
+    messageTreeMapSizeKb: mapSizeKb,
+    publicDataTreeMapSizeKb: mapSizeKb,
+  });
+
+  worldStateCache.set(dataDir, ws);
+  return ws;
+}
+
+async function simulateWithFuzzer(
+  dataDir: string,
+  mapSizeKb: number,
+  txHint: AvmTxHint,
+  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
+  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, protocolContracts);
+
+  await simulator.applyNoteHashes(rawNoteHashes);
+
+  // Register contract classes from C++ (must happen before public data writes to match C++ order)
+  for (const rawClass of rawContractClasses) {
+    await simulator.addContractClassFromCpp(rawClass);
+  }
+
+  // Register contract instances from C++
+  for (const [rawAddress, rawInstance] of rawContractInstances) {
+    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
+    .getAppLogicReturnValues()
+    .flatMap((rv: { values?: Fr[] } | undefined) => rv?.values?.filter((v: Fr | null | undefined) => v != null) ?? []);
+
+  return {
+    reverted: !result.revertCode.isOK(),
+    output,
+    revertReason: result.findRevertReason()?.message,
+    publicInputs: result.publicInputs!,
+    publicTxEffect: result.publicTxEffect,
+  };
+}
+
+async function execute(base64Line: string): Promise<void> {
+  try {
+    // Decode base64 and deserialize the entire request from msgpack
+    const buffer = Buffer.from(base64Line.trim(), 'base64');
+    const rawRequest = deserializeFromMessagePack(buffer);
+    const request = FuzzerSimulationRequest.fromPlainObject(rawRequest);
+
+    // Run the TS simulation
+    const result = await simulateWithFuzzer(
+      request.wsDataDir,
+      request.wsMapSizeKb,
+      request.tx,
+      request.globals,
+      request.contractClasses,
+      request.contractInstances,
+      request.publicDataWrites,
+      request.noteHashes,
+      request.protocolContracts,
+    );
+
+    // Serialize the result to msgpack and encode it in base64 for output
+    const resultBuffer = serializeWithMessagePack({
+      reverted: result.reverted,
+      output: result.output,
+      revertReason: result.revertReason ?? '',
+      endTreeSnapshots: result.publicInputs.endTreeSnapshots,
+      publicTxEffect: result.publicTxEffect,
+    });
+    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 Fr[],
+      revertReason: `Unexpected Error ${error.message}`,
+      endTreeSnapshots: TreeSnapshots.empty(),
+      publicTxEffect: PublicTxEffect.empty(),
+    });
+    await writeOutput(errorResult.toString('base64') + '\n');
+  }
+}
+
+function mainLoop() {
+  const rl = createInterface({ input: process.stdin, terminal: false });
+
+  // Process lines sequentially to avoid race conditions in responses
+  const lineQueue: string[] = [];
+  let processing = false;
+
+  async function processQueue() {
+    if (processing || lineQueue.length === 0) {
+      return;
+    }
+    processing = true;
+    while (lineQueue.length > 0) {
+      const line = lineQueue.shift()!;
+      await execute(line);
+    }
+    processing = false;
+  }
+
+  rl.on('line', (line: string) => {
+    if (line.trim()) {
+      lineQueue.push(line);
+      void processQueue();
+    }
+  });
+  rl.on('close', () => process.exit(0));
+}
+
+void mainLoop();

package/src/public/hinting_db_sources.ts

@@ -1,5 +1,6 @@
-import { sha256Trunc } from '@aztec/foundation/crypto';
-import { Fr } from '@aztec/foundation/fields';
+import type { BlockNumber } from '@aztec/foundation/branded-types';
+import { sha256Trunc } from '@aztec/foundation/crypto/sha256';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import type { IndexedTreeLeafPreimage, SiblingPath } from '@aztec/foundation/trees';
 import type { FunctionSelector } from '@aztec/stdlib/abi';
@@ -409,12 +410,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();
 
@@ -423,14 +424,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> {
@@ -571,6 +574,10 @@ export class HintingMerkleWriteOperations implements MerkleTreeWriteOperations {
     return await this.db.close();
   }
 
+  async [Symbol.asyncDispose](): Promise<void> {
+    await this.close();
+  }
+
   public async findLeafIndices<ID extends MerkleTreeId>(
     treeId: ID,
     values: MerkleTreeLeafType<ID>[],
@@ -596,7 +603,7 @@ export class HintingMerkleWriteOperations implements MerkleTreeWriteOperations {
   public async getBlockNumbersForLeafIndices<ID extends MerkleTreeId>(
     treeId: ID,
     leafIndices: bigint[],
-  ): Promise<(bigint | undefined)[]> {
+  ): Promise<(BlockNumber | undefined)[]> {
     return await this.db.getBlockNumbersForLeafIndices(treeId, leafIndices);
   }
 }

package/src/public/index.ts

@@ -1,6 +1,12 @@
 export { PublicContractsDB } from './public_db_sources.js';
 export { GuardedMerkleTreeOperations } from './public_processor/guarded_merkle_tree.js';
 export { PublicProcessor, PublicProcessorFactory } from './public_processor/public_processor.js';
-export { PublicTxSimulator, TelemetryPublicTxSimulator } from './public_tx_simulator/index.js';
+export {
+  CppPublicTxSimulator,
+  createPublicTxSimulatorForBlockBuilding,
+  DumpingCppPublicTxSimulator,
+  type PublicTxSimulatorInterface,
+  TelemetryCppPublicTxSimulator,
+} from './public_tx_simulator/index.js';
 export type { PublicTxResult, PublicSimulatorConfig as PublicTxSimulatorConfig } from '@aztec/stdlib/avm';
 export { getCallRequestsWithCalldataByPhase } from './utils.js';

package/src/public/public_db_sources.ts

@@ -4,8 +4,8 @@ import {
   NULLIFIER_SUBTREE_HEIGHT,
   PUBLIC_DATA_SUBTREE_HEIGHT,
 } from '@aztec/constants';
-import { Fr } from '@aztec/foundation/fields';
-import { createLogger } from '@aztec/foundation/log';
+import { Fr } from '@aztec/foundation/curves/bn254';
+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,19 @@ 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> {
+  public addContracts(contractDeploymentData: ContractDeploymentData): void {
     const currentState = this.getCurrentState();
 
-    await this.addContractClassesFromEvents(
+    this.addContractClassesFromEvents(
       ContractClassPublishedEvent.extractContractClassEvents(contractDeploymentData.getContractClassLogs()),
       currentState,
     );
@@ -64,10 +69,10 @@ 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.addContracts(contractDeploymentData.getNonRevertibleContractDeploymentData());
+    this.addContracts(contractDeploymentData.getRevertibleContractDeploymentData());
   }
 
   /**
@@ -76,7 +81,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 +174,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 +211,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();

package/src/public/public_processor/guarded_merkle_tree.ts

@@ -1,3 +1,4 @@
+import type { BlockNumber } from '@aztec/foundation/branded-types';
 import { SerialQueue } from '@aztec/foundation/queue';
 import type { IndexedTreeLeafPreimage, SiblingPath } from '@aztec/foundation/trees';
 import type {
@@ -80,6 +81,10 @@ export class GuardedMerkleTreeOperations implements MerkleTreeWriteOperations {
   close(): Promise<void> {
     return this.guardAndPush(() => this.target.close());
   }
+
+  async [Symbol.asyncDispose](): Promise<void> {
+    await this.close();
+  }
   getTreeInfo(treeId: MerkleTreeId): Promise<TreeInfo> {
     return this.guardAndPush(() => this.target.getTreeInfo(treeId));
   }
@@ -126,10 +131,10 @@ export class GuardedMerkleTreeOperations implements MerkleTreeWriteOperations {
   getBlockNumbersForLeafIndices<ID extends MerkleTreeId>(
     treeId: ID,
     leafIndices: bigint[],
-  ): Promise<(bigint | undefined)[]> {
+  ): Promise<(BlockNumber | undefined)[]> {
     return this.guardAndPush(() => this.target.getBlockNumbersForLeafIndices(treeId, leafIndices));
   }
-  createCheckpoint(): Promise<void> {
+  createCheckpoint(): Promise<number> {
     return this.guardAndPush(() => this.target.createCheckpoint());
   }
   commitCheckpoint(): Promise<void> {
@@ -138,11 +143,11 @@ export class GuardedMerkleTreeOperations implements MerkleTreeWriteOperations {
   revertCheckpoint(): Promise<void> {
     return this.guardAndPush(() => this.target.revertCheckpoint());
   }
-  commitAllCheckpoints(): Promise<void> {
-    return this.guardAndPush(() => this.target.commitAllCheckpoints());
+  commitAllCheckpointsTo(depth: number): Promise<void> {
+    return this.guardAndPush(() => this.target.commitAllCheckpointsTo(depth));
   }
-  revertAllCheckpoints(): Promise<void> {
-    return this.guardAndPush(() => this.target.revertAllCheckpoints());
+  revertAllCheckpointsTo(depth: number): Promise<void> {
+    return this.guardAndPush(() => this.target.revertAllCheckpointsTo(depth));
   }
   findSiblingPaths<ID extends MerkleTreeId>(
     treeId: ID,