@aztec/simulator 3.0.0-devnet.6 → 3.0.0-devnet.6-patch.1

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';
@@ -8,8 +9,12 @@ import {
   AvmBytecodeCommitmentHint,
   AvmCommitCheckpointHint,
   AvmContractClassHint,
+  AvmContractDbCommitCheckpointHint,
+  AvmContractDbCreateCheckpointHint,
+  AvmContractDbRevertCheckpointHint,
   AvmContractInstanceHint,
   AvmCreateCheckpointHint,
+  AvmDebugFunctionNameHint,
   type AvmExecutionHints,
   AvmGetLeafPreimageHintNullifierTree,
   AvmGetLeafPreimageHintPublicDataTree,
@@ -50,8 +55,11 @@ import type { PublicContractsDBInterface } from './db_interfaces.js';
  * A public contracts database that forwards requests and collects AVM hints.
  */
 export class HintingPublicContractsDB implements PublicContractsDBInterface {
-  // We deduplicate contract classes because they include the whole bytecode.
-  private contractClassIds: Set<bigint> = new Set();
+  private static readonly log: Logger = createLogger('simulator:hinting-public-contracts-db');
+
+  private checkpointActionCounter: number = 0;
+  private nextCheckpointId: number = 1;
+  private checkpointStack: number[] = [0];
 
   constructor(
     private readonly db: PublicContractsDBInterface,
@@ -64,9 +72,11 @@ export class HintingPublicContractsDB implements PublicContractsDBInterface {
   ): Promise<ContractInstanceWithAddress | undefined> {
     const instance = await this.db.getContractInstance(address, timestamp);
     if (instance) {
-      // We don't need to hint the block number because it doesn't change.
+      const hintKey = this.getHintKey();
+
       this.hints.contractInstances.push(
         new AvmContractInstanceHint(
+          hintKey,
           instance.address,
           instance.salt,
           instance.deployer,
@@ -76,22 +86,31 @@ export class HintingPublicContractsDB implements PublicContractsDBInterface {
           instance.publicKeys,
         ),
       );
+
+      HintingPublicContractsDB.log.trace(
+        `[getContractInstance:${hintKey}] Added contract instance ${instance.address.toString()} to hints.`,
+      );
     }
     return instance;
   }
 
   public async getContractClass(contractClassId: Fr): Promise<ContractClassPublic | undefined> {
     const contractClass = await this.db.getContractClass(contractClassId);
-    if (contractClass && !this.contractClassIds.has(contractClassId.toBigInt())) {
-      this.contractClassIds.add(contractClassId.toBigInt());
+    if (contractClass) {
+      const hintKey = this.getHintKey();
       this.hints.contractClasses.push(
         new AvmContractClassHint(
+          hintKey,
           contractClass.id,
           contractClass.artifactHash,
           contractClass.privateFunctionsRoot,
           contractClass.packedBytecode,
         ),
       );
+
+      HintingPublicContractsDB.log.trace(
+        `[getContractClass:${hintKey}] Added contract class ${contractClassId.toString()} to hints.`,
+      );
     }
     return contractClass;
   }
@@ -99,7 +118,12 @@ export class HintingPublicContractsDB implements PublicContractsDBInterface {
   public async getBytecodeCommitment(contractClassId: Fr): Promise<Fr | undefined> {
     const commitment = await this.db.getBytecodeCommitment(contractClassId);
     if (commitment) {
-      this.hints.bytecodeCommitments.push(new AvmBytecodeCommitmentHint(contractClassId, commitment));
+      const hintKey = this.getHintKey();
+      this.hints.bytecodeCommitments.push(new AvmBytecodeCommitmentHint(hintKey, contractClassId, commitment));
+
+      HintingPublicContractsDB.log.trace(
+        `[getBytecodeCommitment:${hintKey}] Added bytecode commitment ${commitment.toString()} to hints for contract class ${contractClassId.toString()}.`,
+      );
     }
     return commitment;
   }
@@ -108,7 +132,81 @@ export class HintingPublicContractsDB implements PublicContractsDBInterface {
     contractAddress: AztecAddress,
     selector: FunctionSelector,
   ): Promise<string | undefined> {
-    return await this.db.getDebugFunctionName(contractAddress, selector);
+    const name = await this.db.getDebugFunctionName(contractAddress, selector);
+    if (name) {
+      HintingPublicContractsDB.log.debug(
+        `[getDebugFunctionName] Adding debug function name ${name} to hints for contract ${contractAddress.toString()} and selector ${selector.toString()}.`,
+      );
+      // We hint selector as a field to make things way simpler in C++.
+      this.hints.debugFunctionNames.push(new AvmDebugFunctionNameHint(contractAddress, selector.toField(), name));
+    }
+    return name;
+  }
+
+  public createCheckpoint(): void {
+    const hintKey = this.getHintKey();
+    this.checkpointActionCounter++;
+    const oldCheckpointId = this.getCurrentCheckpointId();
+
+    this.db.createCheckpoint();
+
+    const newCheckpointId = this.nextCheckpointId++;
+    this.checkpointStack.push(newCheckpointId);
+
+    this.hints.contractDbCreateCheckpointHints.push(
+      new AvmContractDbCreateCheckpointHint(hintKey, oldCheckpointId, newCheckpointId),
+    );
+
+    HintingPublicContractsDB.log.trace(
+      `[createCheckpoint:${hintKey}] Checkpoint evolved ${oldCheckpointId} -> ${newCheckpointId}.`,
+    );
+  }
+
+  public commitCheckpoint(): void {
+    const hintKey = this.getHintKey();
+    this.checkpointActionCounter++;
+    const oldCheckpointId = this.getCurrentCheckpointId();
+
+    this.db.commitCheckpoint();
+
+    this.checkpointStack.pop();
+    const newCheckpointId = this.getCurrentCheckpointId();
+
+    this.hints.contractDbCommitCheckpointHints.push(
+      new AvmContractDbCommitCheckpointHint(hintKey, oldCheckpointId, newCheckpointId),
+    );
+
+    HintingPublicContractsDB.log.trace(
+      `[commitCheckpoint:${hintKey}] Checkpoint evolved ${oldCheckpointId} -> ${newCheckpointId}.`,
+    );
+  }
+
+  public revertCheckpoint(): void {
+    const hintKey = this.getHintKey();
+    this.checkpointActionCounter++;
+    const oldCheckpointId = this.getCurrentCheckpointId();
+
+    this.db.revertCheckpoint();
+
+    this.checkpointStack.pop();
+    const newCheckpointId = this.getCurrentCheckpointId();
+
+    this.hints.contractDbRevertCheckpointHints.push(
+      new AvmContractDbRevertCheckpointHint(hintKey, oldCheckpointId, newCheckpointId),
+    );
+
+    HintingPublicContractsDB.log.trace(
+      `[revertCheckpoint:${hintKey}] Checkpoint evolved ${oldCheckpointId} -> ${newCheckpointId}.`,
+    );
+  }
+
+  // Private methods.
+  private getHintKey(): number {
+    return this.checkpointActionCounter;
+  }
+
+  private getCurrentCheckpointId(): number {
+    return this.checkpointStack[this.checkpointStack.length - 1];
   }
 }
 
@@ -499,7 +597,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

@@ -2,9 +2,11 @@ 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,
-  type PublicTxSimulatorConfig,
-  TelemetryPublicTxSimulator,
-  type PublicTxResult,
+  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,7 +4,7 @@ import {
   NULLIFIER_SUBTREE_HEIGHT,
   PUBLIC_DATA_SUBTREE_HEIGHT,
 } from '@aztec/constants';
-import { Fr } from '@aztec/foundation/fields';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import { ContractClassPublishedEvent } from '@aztec/protocol-contracts/class-registry';
@@ -13,14 +13,15 @@ import type { FunctionSelector } from '@aztec/stdlib/abi';
 import { PublicDataWrite } from '@aztec/stdlib/avm';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import {
+  AllContractDeploymentData,
   type ContractClassPublic,
   type ContractDataSource,
+  type ContractDeploymentData,
   type ContractInstanceWithAddress,
   computePublicBytecodeCommitment,
 } from '@aztec/stdlib/contract';
 import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
 import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
-import { ContractClassLog, PrivateLog } from '@aztec/stdlib/logs';
 import type { PublicDBAccessStats } from '@aztec/stdlib/stats';
 import {
   MerkleTreeId,
@@ -34,182 +35,89 @@ import type { UInt64 } from '@aztec/stdlib/types';
 
 import { strict as assert } from 'assert';
 
+import { ContractsDbCheckpoint } from './contracts_db_checkpoint.js';
 import type { PublicContractsDBInterface, PublicStateDBInterface } from './db_interfaces.js';
 import { L1ToL2MessageIndexOutOfRangeError, NoteHashIndexOutOfRangeError } from './side_effect_errors.js';
-import { TxContractCache } from './tx_contract_cache.js';
 
 /**
  * Implements the PublicContractsDBInterface using a ContractDataSource.
- * Progressively records contracts in transaction as they are processed in a block.
- * Separates block-level contract information (from processed/included txs) from the
- * current tx's contract information (which may be cleared on tx revert/death).
+ * Uses a stack-based checkpoint model for managing contract state.
  */
 export class PublicContractsDB implements PublicContractsDBInterface {
-  // Two caching layers for contract classes and instances.
-  // Tx-level cache:
-  //   - The current tx's new contract information is cached
-  //     in currentTxNonRevertibleCache and currentTxRevertibleCache.
-  // Block-level cache:
-  //   - Contract information from earlier in the block, usable by later txs.
-  // When a tx succeeds, that tx's caches are merged into the block cache and cleared.
-  private currentTxNonRevertibleCache = new TxContractCache();
-  private currentTxRevertibleCache = new TxContractCache();
-  private blockCache = new TxContractCache();
-  // Separate flat cache for bytecode commitments.
-  private bytecodeCommitmentCache = new Map<string, Fr>();
+  private contractStateStack: ContractsDbCheckpoint[] = [new ContractsDbCheckpoint()];
 
   private log = createLogger('simulator:contracts-data-source');
 
   constructor(private dataSource: ContractDataSource) {}
 
-  /**
-   * Add new contracts from a transaction
-   * @param tx - The transaction to add contracts from.
-   */
-  public async addNewContracts(tx: Tx): Promise<void> {
-    await this.addNonRevertibleContractClasses(tx);
-    await this.addRevertibleContractClasses(tx);
-    this.addNonRevertibleContractInstances(tx);
-    this.addRevertibleContractInstances(tx);
-  }
+  public async addContracts(contractDeploymentData: ContractDeploymentData): Promise<void> {
+    const currentState = this.getCurrentState();
 
-  /**
-   * Add non revertible contracts from a transaction
-   * @param tx - The transaction to add non revertible contracts from.
-   */
-  public async addNewNonRevertibleContracts(tx: Tx) {
-    await this.addNonRevertibleContractClasses(tx);
-    this.addNonRevertibleContractInstances(tx);
-  }
-
-  /**
-   * Add revertible contracts from a transaction
-   * @param tx - The transaction to add revertible contracts from.
-   */
-  public async addNewRevertibleContracts(tx: Tx) {
-    await this.addRevertibleContractClasses(tx);
-    this.addRevertibleContractInstances(tx);
-  }
-
-  /**
-   * Add non-revertible contract classes from a transaction
-   * For private-only txs, this will be all contract classes (found in tx.data.forPublic)
-   * @param tx - The transaction to add non-revertible contract classes from.
-   */
-  private async addNonRevertibleContractClasses(tx: Tx) {
-    const siloedContractClassLogs = tx.data.forPublic
-      ? tx.getSplitContractClassLogs(false /* revertible */)
-      : tx.getContractClassLogs();
-    await this.addContractClassesFromLogs(siloedContractClassLogs, this.currentTxNonRevertibleCache, 'non-revertible');
-  }
+    await this.addContractClassesFromEvents(
+      ContractClassPublishedEvent.extractContractClassEvents(contractDeploymentData.getContractClassLogs()),
+      currentState,
+    );
 
-  /**
-   * Add revertible contract classes from a transaction
-   * None for private-only txs.
-   * @param tx - The transaction to add revertible contract classes from.
-   */
-  private async addRevertibleContractClasses(tx: Tx) {
-    const siloedContractClassLogs = tx.data.forPublic ? tx.getSplitContractClassLogs(true /* revertible */) : [];
-    await this.addContractClassesFromLogs(siloedContractClassLogs, this.currentTxRevertibleCache, 'revertible');
+    this.addContractInstancesFromEvents(
+      ContractInstancePublishedEvent.extractContractInstanceEvents(contractDeploymentData.getPrivateLogs()),
+      currentState,
+    );
   }
 
-  /**
-   * Add non-revertible contract instances from a transaction
-   * For private-only txs, this will be all contract instances (found in tx.data.forRollup)
-   * @param tx - The transaction to add non-revertible contract instances from.
-   */
-  private addNonRevertibleContractInstances(tx: Tx) {
-    const contractInstanceLogs = tx.data.forPublic
-      ? tx.data.forPublic!.nonRevertibleAccumulatedData.privateLogs.filter(l => !l.isEmpty())
-      : tx.data.forRollup!.end.privateLogs.filter(l => !l.isEmpty());
-
-    this.addContractInstancesFromLogs(contractInstanceLogs, this.currentTxNonRevertibleCache, 'non-revertible');
+  public async addNewContracts(tx: Tx): Promise<void> {
+    const contractDeploymentData = AllContractDeploymentData.fromTx(tx);
+    await this.addContracts(contractDeploymentData.getNonRevertibleContractDeploymentData());
+    await this.addContracts(contractDeploymentData.getRevertibleContractDeploymentData());
   }
 
   /**
-   * Add revertible contract instances from a transaction
-   * None for private-only txs.
-   * @param tx - The transaction to add revertible contract instances from.
+   * Creates a new checkpoint, copying the current state for upcoming modifications,
+   * and enabling rollbacks to current state in case of a revert.
    */
-  private addRevertibleContractInstances(tx: Tx) {
-    const contractInstanceLogs = tx.data.forPublic
-      ? tx.data.forPublic!.revertibleAccumulatedData.privateLogs.filter(l => !l.isEmpty())
-      : [];
-
-    this.addContractInstancesFromLogs(contractInstanceLogs, this.currentTxRevertibleCache, 'revertible');
+  public createCheckpoint(): void {
+    const currentState = this.getCurrentState();
+    const newState = currentState.deepCopy();
+    this.contractStateStack.push(newState);
   }
 
   /**
-   * Given a tx's siloed contract class logs, add the contract classes to the cache
-   * @param siloedContractClassLogs - Contract class logs to process
-   * @param cache - The cache to store the contract classes in
-   * @param cacheType - Type of cache (for logging)
+   * Commits the current checkpoint, accepting its state latest.
    */
-  private async addContractClassesFromLogs(
-    siloedContractClassLogs: ContractClassLog[],
-    cache: TxContractCache,
-    cacheType: string,
-  ) {
-    const contractClassEvents = siloedContractClassLogs
-      .filter((log: ContractClassLog) => ContractClassPublishedEvent.isContractClassPublishedEvent(log))
-      .map((log: ContractClassLog) => ContractClassPublishedEvent.fromLog(log));
-
-    // Cache contract classes
-    await Promise.all(
-      contractClassEvents.map(async (event: ContractClassPublishedEvent) => {
-        this.log.debug(`Adding class ${event.contractClassId.toString()} to contract's ${cacheType} tx cache`);
-        const contractClass = await event.toContractClassPublic();
-
-        cache.addClass(event.contractClassId, contractClass);
-      }),
-    );
-  }
-
-  /**
-   * Given a tx's contract instance logs, add the contract instances to the cache
-   * @param contractInstanceLogs - Contract instance logs to process
-   * @param cache - The cache to store the contract instances in
-   * @param cacheType - Type of cache (for logging)
-   */
-  private addContractInstancesFromLogs(contractInstanceLogs: PrivateLog[], cache: TxContractCache, cacheType: string) {
-    const contractInstanceEvents = contractInstanceLogs
-      .filter(log => ContractInstancePublishedEvent.isContractInstancePublishedEvent(log))
-      .map(log => ContractInstancePublishedEvent.fromLog(log));
-
-    // Cache contract instances
-    contractInstanceEvents.forEach(e => {
-      this.log.debug(
-        `Adding instance ${e.address.toString()} with class ${e.contractClassId.toString()} to ${cacheType} tx contract cache`,
-      );
-      cache.addInstance(e.address, e.toContractInstance());
-    });
+  public commitCheckpoint(): void {
+    if (this.contractStateStack.length <= 1) {
+      throw new Error('No checkpoint to commit');
+    }
+    const topState = this.contractStateStack.pop()!;
+    this.contractStateStack[this.contractStateStack.length - 1] = topState;
   }
 
   /**
-   * Clear new contracts from the current tx's cache
+   * Commits the current checkpoint, not erroring if there is no checkpoint
+   * to commit. This is useful to do a sanity commit at the end of tx execution,
+   * doing nothing if the checkpoint was already reverted, but truly committing
+   * otherwise.
    */
-  public clearContractsForTx() {
-    this.currentTxRevertibleCache.clear();
-    this.currentTxRevertibleCache.clear();
-    this.currentTxNonRevertibleCache.clear();
+  public commitCheckpointOkIfNone(): void {
+    if (this.contractStateStack.length <= 1) {
+      return;
+    }
+    const topState = this.contractStateStack.pop()!;
+    this.contractStateStack[this.contractStateStack.length - 1] = topState;
   }
 
   /**
-   * Commits the current transaction's cached contracts to the block-level cache.
-   * Then, clears the tx cache.
+   * Reverts the current checkpoint, discarding its state and rolling back
+   * to the state as of the latest checkpoint.
    */
-  public commitContractsForTx(onlyNonRevertibles: boolean = false) {
-    // Merge non-revertible tx cache into block cache
-    this.blockCache.mergeFrom(this.currentTxNonRevertibleCache);
-
-    if (!onlyNonRevertibles) {
-      // Merge revertible tx cache into block cache
-      this.blockCache.mergeFrom(this.currentTxRevertibleCache);
+  public revertCheckpoint(): void {
+    if (this.contractStateStack.length <= 1) {
+      throw new Error('No checkpoint to revert');
     }
+    this.contractStateStack.pop();
+  }
 
-    // Clear the tx's caches
-    this.currentTxNonRevertibleCache.clear();
-    this.currentTxRevertibleCache.clear();
+  private getCurrentState(): ContractsDbCheckpoint {
+    return this.contractStateStack[this.contractStateStack.length - 1];
   }
 
   // TODO(fcarreiro/alvaro): This method currently needs a blockNumber. Since this class
@@ -225,53 +133,66 @@ export class PublicContractsDB implements PublicContractsDBInterface {
     address: AztecAddress,
     timestamp: UInt64,
   ): Promise<ContractInstanceWithAddress | undefined> {
-    // Check caches in order: tx revertible -> tx non-revertible -> block -> data source
-    return (
-      this.currentTxRevertibleCache.getInstance(address) ??
-      this.currentTxNonRevertibleCache.getInstance(address) ??
-      this.blockCache.getInstance(address) ??
-      (await this.dataSource.getContract(address, timestamp))
-    );
+    const currentState = this.getCurrentState();
+    return currentState.getInstance(address) ?? (await this.dataSource.getContract(address, timestamp));
   }
 
   public async getContractClass(contractClassId: Fr): Promise<ContractClassPublic | undefined> {
-    // Check caches in order: tx revertible -> tx non-revertible -> block -> data source
-    return (
-      this.currentTxRevertibleCache.getClass(contractClassId) ??
-      this.currentTxNonRevertibleCache.getClass(contractClassId) ??
-      this.blockCache.getClass(contractClassId) ??
-      (await this.dataSource.getContractClass(contractClassId))
-    );
+    const currentState = this.getCurrentState();
+    return currentState.getClass(contractClassId) ?? (await this.dataSource.getContractClass(contractClassId));
   }
 
   public async getBytecodeCommitment(contractClassId: Fr): Promise<Fr | undefined> {
-    // Try and retrieve from cache
-    const key = contractClassId.toString();
-    const result = this.bytecodeCommitmentCache.get(key);
-    if (result !== undefined) {
-      return result;
+    const currentState = this.getCurrentState();
+    const commitment =
+      currentState.getBytecodeCommitment(contractClassId) ??
+      (await this.dataSource.getBytecodeCommitment(contractClassId));
+    if (commitment !== undefined) {
+      return commitment;
     }
-    // Now try from the store
-    const fromStore = await this.dataSource.getBytecodeCommitment(contractClassId);
-    if (fromStore !== undefined) {
-      this.bytecodeCommitmentCache.set(key, fromStore);
-      return fromStore;
-    }
-
-    // Not in either the store or the cache, build it here and cache
+    // Not in the current state or the store, compute it here
+    // Get the contract class
     const contractClass = await this.getContractClass(contractClassId);
+
     if (contractClass === undefined) {
+      // cannot compute bytecode commitment if contract class is not found
       return undefined;
     }
 
     const value = await computePublicBytecodeCommitment(contractClass.packedBytecode);
-    this.bytecodeCommitmentCache.set(key, value);
+    // Add to cache (current checkpoint state) so we don't compute again
+    currentState.addBytecodeCommitment(contractClassId, value);
     return value;
   }
 
   public async getDebugFunctionName(address: AztecAddress, selector: FunctionSelector): Promise<string | undefined> {
     return await this.dataSource.getDebugFunctionName(address, selector);
   }
+
+  private async 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);
+      }),
+    );
+  }
+
+  private addContractInstancesFromEvents(
+    contractInstanceEvents: ContractInstancePublishedEvent[],
+    state: ContractsDbCheckpoint,
+  ) {
+    contractInstanceEvents.forEach(e => {
+      this.log.debug(
+        `Adding instance ${e.address.toString()} with class ${e.contractClassId.toString()} to contract state`,
+      );
+      state.addInstance(e.address, e.toContractInstance());
+    });
+  }
 }
 
 /**

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 {
@@ -11,7 +12,7 @@ import type {
   TreeInfo,
 } from '@aztec/stdlib/trees';
 import type { BlockHeader, StateReference } from '@aztec/stdlib/tx';
-import type { WorldStateRevision } from '@aztec/stdlib/world-state';
+import type { WorldStateRevision, WorldStateRevisionWithHandle } from '@aztec/stdlib/world-state';
 
 /**
  * Wraps an instance of `MerkleTreeWriteOperations` to allow the sequencer to gate access.
@@ -89,7 +90,7 @@ export class GuardedMerkleTreeOperations implements MerkleTreeWriteOperations {
   getInitialHeader(): BlockHeader {
     return this.target.getInitialHeader();
   }
-  public getRevision(): WorldStateRevision {
+  public getRevision(): WorldStateRevision | WorldStateRevisionWithHandle {
     return this.target.getRevision();
   }
   getSiblingPath<ID extends MerkleTreeId>(treeId: ID, index: bigint): Promise<SiblingPath<TreeHeights[ID]>> {
@@ -126,7 +127,7 @@ 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> {