@aztec/simulator 0.0.0-test.1 → 0.0.1-commit.0208eb9

package/src/public/public_db_sources.ts

@@ -1,327 +1,278 @@
-import { Fr } from '@aztec/foundation/fields';
-import { createLogger } from '@aztec/foundation/log';
+import {
+  L1_TO_L2_MSG_TREE_LEAF_COUNT,
+  NOTE_HASH_TREE_LEAF_COUNT,
+  NULLIFIER_SUBTREE_HEIGHT,
+  PUBLIC_DATA_SUBTREE_HEIGHT,
+} from '@aztec/constants';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
-import { ContractClassRegisteredEvent } from '@aztec/protocol-contracts/class-registerer';
-import { ContractInstanceDeployedEvent } from '@aztec/protocol-contracts/instance-deployer';
+import { ContractClassPublishedEvent } from '@aztec/protocol-contracts/class-registry';
+import { ContractInstancePublishedEvent } from '@aztec/protocol-contracts/instance-registry';
 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 {
-  MerkleTreeCheckpointOperations,
-  MerkleTreeReadOperations,
-  MerkleTreeWriteOperations,
-} from '@aztec/stdlib/interfaces/server';
-import { ContractClassLog, PrivateLog } from '@aztec/stdlib/logs';
+import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
 import type { PublicDBAccessStats } from '@aztec/stdlib/stats';
-import { MerkleTreeId, type PublicDataTreeLeafPreimage } from '@aztec/stdlib/trees';
-import type { Tx } from '@aztec/stdlib/tx';
+import {
+  MerkleTreeId,
+  NullifierLeaf,
+  PublicDataTreeLeaf,
+  type PublicDataTreeLeafPreimage,
+  getTreeName,
+} from '@aztec/stdlib/trees';
+import { TreeSnapshots, type Tx } from '@aztec/stdlib/tx';
+import type { UInt64 } from '@aztec/stdlib/types';
+
+import { strict as assert } from 'assert';
 
-import type { PublicContractsDB, PublicStateDB } from '../common/db_interfaces.js';
-import { TxContractCache } from './tx_contract_cache.js';
+import { ContractsDbCheckpoint } from './contracts_db_checkpoint.js';
+import type { PublicContractsDBInterface, PublicStateDBInterface } from './db_interfaces.js';
+import { L1ToL2MessageIndexOutOfRangeError, NoteHashIndexOutOfRangeError } from './side_effect_errors.js';
 
 /**
- * Implements the PublicContractsDB 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).
+ * Implements the PublicContractsDBInterface using a ContractDataSource.
+ * Uses a stack-based checkpoint model for managing contract state.
  */
-export class ContractsDataSourcePublicDB implements PublicContractsDB {
-  // 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 log = createLogger('simulator:contracts-data-source');
-
-  constructor(private dataSource: ContractDataSource) {}
+export class PublicContractsDB implements PublicContractsDBInterface {
+  private contractStateStack: ContractsDbCheckpoint[] = [new ContractsDbCheckpoint()];
 
-  /**
-   * 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);
-  }
+  private log: Logger;
 
-  /**
-   * 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
-      ? await tx.filterContractClassLogs(
-          tx.data.forPublic!.nonRevertibleAccumulatedData.contractClassLogsHashes,
-          /*siloed=*/ true,
-        )
-      : await tx.filterContractClassLogs(tx.data.forRollup!.end.contractClassLogsHashes, /*siloed=*/ true);
-
-    await this.addContractClassesFromLogs(siloedContractClassLogs, this.currentTxNonRevertibleCache, 'non-revertible');
+  constructor(
+    private dataSource: ContractDataSource,
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('simulator:contracts-data-source', bindings);
   }
 
-  /**
-   * 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
-      ? await tx.filterContractClassLogs(
-          tx.data.forPublic!.revertibleAccumulatedData.contractClassLogsHashes,
-          /*siloed=*/ true,
-        )
-      : [];
-
-    await this.addContractClassesFromLogs(siloedContractClassLogs, this.currentTxRevertibleCache, 'revertible');
-  }
+  public async addContracts(contractDeploymentData: ContractDeploymentData): Promise<void> {
+    const currentState = this.getCurrentState();
 
-  /**
-   * 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());
+    await this.addContractClassesFromEvents(
+      ContractClassPublishedEvent.extractContractClassEvents(contractDeploymentData.getContractClassLogs()),
+      currentState,
+    );
 
-    this.addContractInstancesFromLogs(contractInstanceLogs, this.currentTxNonRevertibleCache, 'non-revertible');
+    this.addContractInstancesFromEvents(
+      ContractInstancePublishedEvent.extractContractInstanceEvents(contractDeploymentData.getPrivateLogs()),
+      currentState,
+    );
   }
 
-  /**
-   * Add revertible contract instances from a transaction
-   * None for private-only txs.
-   * @param tx - The transaction to add revertible contract instances from.
-   */
-  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 async addNewContracts(tx: Tx): Promise<void> {
+    const contractDeploymentData = AllContractDeploymentData.fromTx(tx);
+    await this.addContracts(contractDeploymentData.getNonRevertibleContractDeploymentData());
+    await this.addContracts(contractDeploymentData.getRevertibleContractDeploymentData());
   }
 
   /**
-   * 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)
+   * Creates a new checkpoint, copying the current state for upcoming modifications,
+   * and enabling rollbacks to current state in case of a revert.
    */
-  private async addContractClassesFromLogs(
-    siloedContractClassLogs: ContractClassLog[],
-    cache: TxContractCache,
-    cacheType: string,
-  ) {
-    const contractClassEvents = siloedContractClassLogs
-      .filter((log: ContractClassLog) => ContractClassRegisteredEvent.isContractClassRegisteredEvent(log))
-      .map((log: ContractClassLog) => ContractClassRegisteredEvent.fromLog(log));
-
-    // Cache contract classes
-    await Promise.all(
-      contractClassEvents.map(async (event: ContractClassRegisteredEvent) => {
-        this.log.debug(`Adding class ${event.contractClassId.toString()} to contract's ${cacheType} tx cache`);
-        const contractClass = await event.toContractClassPublic();
-
-        cache.addClass(event.contractClassId, contractClass);
-      }),
-    );
+  public createCheckpoint(): void {
+    const currentState = this.getCurrentState();
+    const newState = currentState.deepCopy();
+    this.contractStateStack.push(newState);
   }
 
   /**
-   * 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)
+   * Commits the current checkpoint, accepting its state latest.
    */
-  private addContractInstancesFromLogs(contractInstanceLogs: PrivateLog[], cache: TxContractCache, cacheType: string) {
-    const contractInstanceEvents = contractInstanceLogs
-      .filter(log => ContractInstanceDeployedEvent.isContractInstanceDeployedEvent(log))
-      .map(log => ContractInstanceDeployedEvent.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];
   }
 
-  public async getContractInstance(address: AztecAddress): 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))
-    );
+  // TODO(fcarreiro/alvaro): This method currently needs a blockNumber. Since this class
+  // is only ever used for a given block, it should be possible to construct it with the
+  // block number and then forget about it. However, since this class (and interface) is
+  // currently more externally exposed than we'd want to, Facundo preferred to not add it
+  // to the constructor right now. If we can make this class more private, we should
+  // reconsider this. A litmus test is in how many places we need to initialize with a
+  // dummy block number (tests or not) and pass block numbers to `super`.
+  // Note: Block number got changed to timestamp so this comment ^ is outdated. Keeping
+  // the comment as is as I am not part of the AVM cabal.
+  public async getContractInstance(
+    address: AztecAddress,
+    timestamp: UInt64,
+  ): Promise<ContractInstanceWithAddress | undefined> {
+    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;
-    }
-    // Now try from the store
-    const fromStore = await this.dataSource.getBytecodeCommitment(contractClassId);
-    if (fromStore !== undefined) {
-      this.bytecodeCommitmentCache.set(key, fromStore);
-      return fromStore;
+    const currentState = this.getCurrentState();
+    const commitment =
+      currentState.getBytecodeCommitment(contractClassId) ??
+      (await this.dataSource.getBytecodeCommitment(contractClassId));
+    if (commitment !== undefined) {
+      return commitment;
     }
-
-    // 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.getContractFunctionName(address, selector);
+    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());
+    });
   }
 }
 
 /**
- * A public state DB that reads and writes to the world state.
+ * A high-level class that provides access to the merkle trees.
+ *
+ * This class is just a helper wrapper around a merkle db. Anything that you can do with it
+ * can also be done directly with the merkle db. This class should NOT be exposed or used
+ * outside of `simulator/src/public`.
+ *
+ * NOTE: This class is currently written in such a way that it would generate the
+ * necessary hints if used with a hinting merkle db. This is a bit of a leak of concepts.
+ * Eventually we can have everything depend on a config/factory at the TxSimulator level
+ * to decide whether to use hints or not (same with tracing, etc).
  */
-export class WorldStateDB extends ContractsDataSourcePublicDB implements PublicStateDB, MerkleTreeCheckpointOperations {
-  private logger = createLogger('simulator:world-state-db');
+export class PublicTreesDB implements PublicStateDBInterface {
+  private logger: Logger;
 
-  constructor(public db: MerkleTreeWriteOperations, dataSource: ContractDataSource) {
-    super(dataSource);
+  constructor(
+    private readonly db: MerkleTreeWriteOperations,
+    bindings?: LoggerBindings,
+  ) {
+    this.logger = createLogger('simulator:public-trees-db', bindings);
   }
 
-  /**
-   * Checkpoints the current fork state
-   */
-  public async createCheckpoint() {
-    await this.db.createCheckpoint();
-  }
+  public async storageRead(contract: AztecAddress, slot: Fr): Promise<Fr> {
+    const timer = new Timer();
+    const leafSlot = (await computePublicDataTreeLeafSlot(contract, slot)).toBigInt();
 
-  /**
-   * Commits the current checkpoint
-   */
-  public async commitCheckpoint() {
-    await this.db.commitCheckpoint();
-  }
+    const lowLeafResult = await this.db.getPreviousValueIndex(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot);
+    if (!lowLeafResult) {
+      throw new Error('Low leaf not found');
+    }
 
-  /**
-   * Reverts the current checkpoint
-   */
-  public async revertCheckpoint() {
-    await this.db.revertCheckpoint();
-  }
+    // TODO: We need this for the hints. See class comment for more details.
+    await this.db.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, lowLeafResult.index);
+    // Unconditionally fetching the preimage for the hints. Move it to the hinting layer?
+    const preimage = (await this.db.getLeafPreimage(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      lowLeafResult.index,
+    )) as PublicDataTreeLeafPreimage;
 
-  public getMerkleInterface(): MerkleTreeWriteOperations {
-    return this.db;
-  }
+    const result = lowLeafResult.alreadyPresent ? preimage.leaf.value : Fr.ZERO;
+    this.logger.debug(`Storage read (contract=${contract}, slot=${slot}, value=${result})`, {
+      eventName: 'public-db-access',
+      duration: timer.ms(),
+      operation: 'storage-read',
+    } satisfies PublicDBAccessStats);
 
-  /**
-   * Reads a value from public storage, returning zero if none.
-   * @param contract - Owner of the storage.
-   * @param slot - Slot to read in the contract storage.
-   * @returns The current value in the storage slot.
-   */
-  public async storageRead(contract: AztecAddress, slot: Fr): Promise<Fr> {
-    return await readPublicState(this.db, contract, slot);
+    return result;
   }
 
-  /**
-   * Records a write to public storage.
-   * @param contract - Owner of the storage.
-   * @param slot - Slot to read in the contract storage.
-   * @param newValue - The new value to store.
-   * @returns The slot of the written leaf in the public data tree.
-   */
   public async storageWrite(contract: AztecAddress, slot: Fr, newValue: Fr): Promise<void> {
+    const timer = new Timer();
     const leafSlot = await computePublicDataTreeLeafSlot(contract, slot);
     const publicDataWrite = new PublicDataWrite(leafSlot, newValue);
     await this.db.sequentialInsert(MerkleTreeId.PUBLIC_DATA_TREE, [publicDataWrite.toBuffer()]);
+
+    this.logger.debug(`Storage write (contract=${contract}, slot=${slot}, value=${newValue})`, {
+      eventName: 'public-db-access',
+      duration: timer.ms(),
+      operation: 'storage-write',
+    } satisfies PublicDBAccessStats);
   }
 
-  public async getL1ToL2LeafValue(leafIndex: bigint): Promise<Fr | undefined> {
+  public async getL1ToL2LeafValue(leafIndex: bigint): Promise<Fr> {
     const timer = new Timer();
+    if (leafIndex > L1_TO_L2_MSG_TREE_LEAF_COUNT) {
+      throw new L1ToL2MessageIndexOutOfRangeError(Number(leafIndex));
+    }
     const leafValue = await this.db.getLeafValue(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, leafIndex);
-    this.logger.debug(`[DB] Fetched L1 to L2 message leaf value`, {
+    assert(leafValue !== undefined, 'Unexpected null response from l1 to l2 message tree');
+    // TODO: We need this for the hints. See class comment for more details.
+    await this.db.getSiblingPath(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, leafIndex);
+
+    this.logger.debug(`Fetched L1 to L2 message leaf value (leafIndex=${leafIndex}, value=${leafValue})`, {
       eventName: 'public-db-access',
       duration: timer.ms(),
       operation: 'get-l1-to-l2-message-leaf-value',
@@ -329,41 +280,120 @@ export class WorldStateDB extends ContractsDataSourcePublicDB implements PublicS
     return leafValue;
   }
 
-  public async getCommitmentValue(leafIndex: bigint): Promise<Fr | undefined> {
+  public async getNoteHash(leafIndex: bigint): Promise<Fr> {
     const timer = new Timer();
+    if (leafIndex > NOTE_HASH_TREE_LEAF_COUNT) {
+      throw new NoteHashIndexOutOfRangeError(Number(leafIndex));
+    }
     const leafValue = await this.db.getLeafValue(MerkleTreeId.NOTE_HASH_TREE, leafIndex);
-    this.logger.debug(`[DB] Fetched commitment leaf value`, {
+    assert(leafValue !== undefined, 'Unexpected null response from note hash tree');
+    // TODO: We need this for the hints. See class comment for more details.
+    await this.db.getSiblingPath(MerkleTreeId.NOTE_HASH_TREE, leafIndex);
+
+    this.logger.debug(`Fetched note hash leaf value (leafIndex=${leafIndex}, value=${leafValue})`, {
       eventName: 'public-db-access',
       duration: timer.ms(),
-      operation: 'get-commitment-leaf-value',
+      operation: 'get-note-hash',
     } satisfies PublicDBAccessStats);
     return leafValue;
   }
 
-  public async getNullifierIndex(nullifier: Fr): Promise<bigint | undefined> {
+  public async writeNoteHash(noteHash: Fr): Promise<void> {
     const timer = new Timer();
-    const index = (await this.db.findLeafIndices(MerkleTreeId.NULLIFIER_TREE, [nullifier.toBuffer()]))[0];
-    this.logger.debug(`[DB] Fetched nullifier index`, {
+    await this.db.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, [noteHash]);
+
+    this.logger.debug(`Wrote note hash (noteHash=${noteHash})`, {
       eventName: 'public-db-access',
       duration: timer.ms(),
-      operation: 'get-nullifier-index',
+      operation: 'write-note-hash',
+    } satisfies PublicDBAccessStats);
+  }
+
+  public async checkNullifierExists(nullifier: Fr): Promise<boolean> {
+    const timer = new Timer();
+    const lowLeafResult = await this.db.getPreviousValueIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBigInt());
+    if (!lowLeafResult) {
+      throw new Error('Low leaf not found');
+    }
+    // TODO: We need this for the hints. See class comment for more details.
+    await this.db.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, lowLeafResult.index);
+    // TODO: We need this for the hints. See class comment for more details.
+    await this.db.getLeafPreimage(MerkleTreeId.NULLIFIER_TREE, lowLeafResult.index);
+    const exists = lowLeafResult.alreadyPresent;
+
+    this.logger.debug(`Checked nullifier exists (nullifier=${nullifier}, exists=${exists})`, {
+      eventName: 'public-db-access',
+      duration: timer.ms(),
+      operation: 'check-nullifier-exists',
+    } satisfies PublicDBAccessStats);
+    return exists;
+  }
+
+  public async writeNullifier(siloedNullifier: Fr): Promise<void> {
+    const timer = new Timer();
+    await this.db.sequentialInsert(MerkleTreeId.NULLIFIER_TREE, [siloedNullifier.toBuffer()]);
+
+    this.logger.debug(`Wrote nullifier (nullifier=${siloedNullifier})`, {
+      eventName: 'public-db-access',
+      duration: timer.ms(),
+      operation: 'write-nullifier',
     } satisfies PublicDBAccessStats);
-    return index;
   }
-}
 
-export async function readPublicState(db: MerkleTreeReadOperations, contract: AztecAddress, slot: Fr): Promise<Fr> {
-  const leafSlot = (await computePublicDataTreeLeafSlot(contract, slot)).toBigInt();
+  public async padTree(treeId: MerkleTreeId, leavesToInsert: number): Promise<void> {
+    const timer = new Timer();
+
+    switch (treeId) {
+      // Indexed trees.
+      case MerkleTreeId.NULLIFIER_TREE:
+        await this.db.batchInsert(
+          treeId,
+          Array(leavesToInsert).fill(NullifierLeaf.empty().toBuffer()),
+          NULLIFIER_SUBTREE_HEIGHT,
+        );
+        break;
+      case MerkleTreeId.PUBLIC_DATA_TREE:
+        await this.db.batchInsert(
+          treeId,
+          Array(leavesToInsert).fill(PublicDataTreeLeaf.empty().toBuffer()),
+          PUBLIC_DATA_SUBTREE_HEIGHT,
+        );
+        break;
+      // Append-only trees.
+      case MerkleTreeId.L1_TO_L2_MESSAGE_TREE:
+      case MerkleTreeId.NOTE_HASH_TREE:
+        await this.db.appendLeaves(treeId, Array(leavesToInsert).fill(Fr.ZERO));
+        break;
+      default:
+        throw new Error(`Padding not supported for tree ${treeId}`);
+    }
+
+    this.logger.debug(`Padded tree (tree=${getTreeName(treeId)}, leavesToInsert=${leavesToInsert})`, {
+      eventName: 'public-db-access',
+      duration: timer.ms(),
+      operation: 'pad-tree',
+    } satisfies PublicDBAccessStats);
+  }
 
-  const lowLeafResult = await db.getPreviousValueIndex(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot);
-  if (!lowLeafResult || !lowLeafResult.alreadyPresent) {
-    return Fr.ZERO;
+  public async createCheckpoint(): Promise<void> {
+    await this.db.createCheckpoint();
+  }
+
+  public async commitCheckpoint(): Promise<void> {
+    await this.db.commitCheckpoint();
   }
 
-  const preimage = (await db.getLeafPreimage(
-    MerkleTreeId.PUBLIC_DATA_TREE,
-    lowLeafResult.index,
-  )) as PublicDataTreeLeafPreimage;
+  public async revertCheckpoint(): Promise<void> {
+    await this.db.revertCheckpoint();
+  }
 
-  return preimage.value;
+  public async getTreeSnapshots(): Promise<TreeSnapshots> {
+    const stateReference = await this.db.getStateReference();
+    return new TreeSnapshots(
+      stateReference.l1ToL2MessageTree,
+      stateReference.partial.noteHashTree,
+      stateReference.partial.nullifierTree,
+      stateReference.partial.publicDataTree,
+    );
+  }
 }