@aztec/pxe 4.0.0-nightly.20260111 → 4.0.0-nightly.20260113

package/src/private_kernel/private_kernel_execution_prover.ts

@@ -28,7 +28,6 @@ import {
   type PrivateCallExecutionResult,
   type PrivateExecutionResult,
   TxRequest,
-  collectNoteHashLeafIndexMap,
   collectNoteHashNullifierCounterMap,
   getFinalMinRevertibleSideEffectCounter,
 } from '@aztec/stdlib/tx';
@@ -101,7 +100,6 @@ export class PrivateKernelExecutionProver {
 
     const executionSteps: PrivateExecutionStep[] = [];
 
-    const noteHashLeafIndexMap = collectNoteHashLeafIndexMap(executionResult);
     const noteHashNullifierCounterMap = collectNoteHashNullifierCounterMap(executionResult);
     const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(executionResult);
     const splitCounter = isPrivateOnlyTx ? 0 : minRevertibleSideEffectCounter;
@@ -116,7 +114,7 @@ export class PrivateKernelExecutionProver {
         );
         while (resetBuilder.needsReset()) {
           const witgenTimer = new Timer();
-          const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
+          const privateInputs = await resetBuilder.build(this.oracle);
           output = generateWitnesses
             ? await this.proofCreator.generateResetOutput(privateInputs)
             : await this.proofCreator.simulateReset(privateInputs);
@@ -224,7 +222,7 @@ export class PrivateKernelExecutionProver {
     );
     while (resetBuilder.needsReset()) {
       const witgenTimer = new Timer();
-      const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
+      const privateInputs = await resetBuilder.build(this.oracle);
       output = generateWitnesses
         ? await this.proofCreator.generateResetOutput(privateInputs)
         : await this.proofCreator.simulateReset(privateInputs);

package/src/private_kernel/private_kernel_oracle.ts

@@ -43,7 +43,8 @@ export interface PrivateKernelOracle {
 
   /**
    * Returns a membership witness with the sibling path and leaf index in our private function indexed merkle tree.
-   */ getNoteHashMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>>;
+   */
+  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined>;
 
   /**
    * Returns a membership witness with the sibling path and leaf index in our nullifier indexed merkle tree.

package/src/private_kernel/private_kernel_oracle_impl.ts

@@ -2,7 +2,6 @@ import { NOTE_HASH_TREE_HEIGHT, PUBLIC_DATA_TREE_HEIGHT, VK_TREE_HEIGHT } from '
 import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { GrumpkinScalar, Point } from '@aztec/foundation/curves/grumpkin';
 import { createLogger } from '@aztec/foundation/log';
-import type { Tuple } from '@aztec/foundation/serialize';
 import { MembershipWitness } from '@aztec/foundation/trees';
 import type { KeyStore } from '@aztec/key-store';
 import { getVKIndex, getVKSiblingPath } from '@aztec/noir-protocol-circuits-types/vk-tree';
@@ -70,13 +69,8 @@ export class PrivateKernelOracleImpl implements PrivateKernelOracle {
     return Promise.resolve(new MembershipWitness(VK_TREE_HEIGHT, BigInt(leafIndex), getVKSiblingPath(leafIndex)));
   }
 
-  async getNoteHashMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>> {
-    const path = await this.node.getNoteHashSiblingPath(this.blockNumber, leafIndex);
-    return new MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>(
-      path.pathSize,
-      leafIndex,
-      path.toFields() as Tuple<Fr, typeof NOTE_HASH_TREE_HEIGHT>,
-    );
+  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined> {
+    return this.node.getNoteHashMembershipWitness(this.blockNumber, noteHash);
   }
 
   getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitness | undefined> {

package/src/pxe.ts

@@ -148,6 +148,7 @@ export class PXE {
     const tipsStore = new L2TipsKVStore(store, 'pxe');
     const synchronizer = new BlockSynchronizer(
       node,
+      store,
       anchorBlockStore,
       noteStore,
       privateEventStore,
@@ -157,6 +158,7 @@ export class PXE {
     );
 
     const jobCoordinator = new JobCoordinator(store);
+    jobCoordinator.registerStores([capsuleStore]);
 
     const debugUtils = new PXEDebugUtils(contractStore, noteStore);
 

package/src/storage/capsule_store/capsule_store.ts

@@ -3,12 +3,21 @@ import { type Logger, createLogger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 
-export class CapsuleStore {
+import type { StagedStore } from '../../job_coordinator/job_coordinator.js';
+
+export class CapsuleStore implements StagedStore {
+  readonly storeName = 'capsule';
+
   #store: AztecAsyncKVStore;
 
   // Arbitrary data stored by contracts. Key is computed as `${contractAddress}:${key}`
   #capsules: AztecAsyncMap<string, Buffer>;
 
+  // jobId => `${contractAddress}:${key}` => capsule data
+  // when `#stagedCapsules.get('some-job-id').get('${some-contract-address:some-key') === null`,
+  // it signals that the capsule was deleted during the job, so it needs to be deleted on commit
+  #stagedCapsules: Map<string, Map<string, Buffer | null>>;
+
   logger: Logger;
 
   constructor(store: AztecAsyncKVStore) {
@@ -16,21 +25,120 @@ export class CapsuleStore {
 
     this.#capsules = this.#store.openMap('capsules');
 
+    this.#stagedCapsules = new Map();
+
     this.logger = createLogger('pxe:capsule-data-provider');
   }
 
+  /**
+   * Given a job denoted by `jobId`, it returns the
+   * capsules that said job has interacted with.
+   *
+   * Capsules that haven't been committed to persistence KV storage
+   * are kept in-memory in `#stagedCapsules`, this method provides a convenient
+   * way to access that in-memory collection of data.
+   *
+   * @param jobId
+   * @returns
+   */
+  #getJobStagedCapsules(jobId: string): Map<string, Buffer | null> {
+    let jobStagedCapsules = this.#stagedCapsules.get(jobId);
+    if (!jobStagedCapsules) {
+      jobStagedCapsules = new Map();
+      this.#stagedCapsules.set(jobId, jobStagedCapsules);
+    }
+    return jobStagedCapsules;
+  }
+
+  /**
+   * Reads a capsule's slot from the staged version of the data associated to the given jobId.
+   *
+   * If it is not there, it reads it from the KV store.
+   */
+  async #getFromStage(jobId: string, dbSlotKey: string): Promise<Buffer | null | undefined> {
+    const jobStagedCapsules = this.#getJobStagedCapsules(jobId);
+    let staged: Buffer | null | undefined = jobStagedCapsules.get(dbSlotKey);
+    // Note that if staged === null, we marked it for deletion, so we don't want to
+    // re-read it from DB
+    if (staged === undefined) {
+      // If we don't have a staged version of this dbSlotKey, first we check if there's one in DB
+      staged = await this.#loadCapsuleFromDb(dbSlotKey);
+    }
+    return staged;
+  }
+
+  /**
+   * Writes a capsule to the stage of a job.
+   */
+  #setOnStage(jobId: string, dbSlotKey: string, capsuleData: Buffer) {
+    this.#getJobStagedCapsules(jobId).set(dbSlotKey, capsuleData);
+  }
+
+  /**
+   * Deletes a capsule on the stage of a job. Note the capsule will still
+   * exist in storage until the job is committed.
+   */
+  #deleteOnStage(jobId: string, dbSlotKey: string) {
+    this.#getJobStagedCapsules(jobId).set(dbSlotKey, null);
+  }
+
+  async #loadCapsuleFromDb(dbSlotKey: string): Promise<Buffer | null> {
+    const dataBuffer = await this.#capsules.getAsync(dbSlotKey);
+    if (!dataBuffer) {
+      return null;
+    }
+
+    return dataBuffer;
+  }
+
+  /**
+   * Commits staged data to main storage.
+   * Called by JobCoordinator when a job completes successfully.
+   * Note: JobCoordinator wraps all commits in a single transaction, so we don't
+   * need our own transactionAsync here (and using one would deadlock on IndexedDB).
+   * @param jobId - The jobId identifying which staged data to commit
+   */
+  async commit(jobId: string): Promise<void> {
+    const jobStagedCapsules = this.#getJobStagedCapsules(jobId);
+
+    for (const [key, value] of jobStagedCapsules) {
+      // In the write stage, we represent deleted capsules with null
+      // (as opposed to undefined, which denotes there was never a capsule there to begin with).
+      // So we delete from actual KV store here.
+      if (value === null) {
+        await this.#capsules.delete(key);
+      } else {
+        await this.#capsules.set(key, value);
+      }
+    }
+
+    this.#stagedCapsules.delete(jobId);
+  }
+
+  /**
+   * Discards staged data without committing.
+   */
+  discardStaged(jobId: string): Promise<void> {
+    this.#stagedCapsules.delete(jobId);
+    return Promise.resolve();
+  }
+
   /**
    * Stores arbitrary information in a per-contract non-volatile database, which can later be retrieved with `loadCapsule`.
    * * If data was already stored at this slot, it is overwritten.
    * @param contractAddress - The contract address to scope the data under.
    * @param slot - The slot in the database in which to store the value. Slots need not be contiguous.
    * @param capsule - An array of field elements representing the capsule.
+   * @param jobId - The context in which this store will be visible until PXE decides to persist it to underlying KV store
    * @remarks A capsule is a "blob" of data that is passed to the contract through an oracle. It works similarly
    * to public contract storage in that it's indexed by the contract address and storage slot but instead of the global
    * network state it's backed by local PXE db.
    */
-  async storeCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[]): Promise<void> {
-    await this.#capsules.set(dbSlotToKey(contractAddress, slot), Buffer.concat(capsule.map(value => value.toBuffer())));
+  storeCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[], jobId: string) {
+    const dbSlotKey = dbSlotToKey(contractAddress, slot);
+
+    // A store overrides any pre-existing data on the slot
+    this.#setOnStage(jobId, dbSlotKey, Buffer.concat(capsule.map(value => value.toBuffer())));
   }
 
   /**
@@ -39,8 +147,8 @@ export class CapsuleStore {
    * @param slot - The slot in the database to read.
    * @returns The stored data or `null` if no data is stored under the slot.
    */
-  async loadCapsule(contractAddress: AztecAddress, slot: Fr): Promise<Fr[] | null> {
-    const dataBuffer = await this.#capsules.getAsync(dbSlotToKey(contractAddress, slot));
+  async loadCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string): Promise<Fr[] | null> {
+    const dataBuffer = await this.#getFromStage(jobId, dbSlotToKey(contractAddress, slot));
     if (!dataBuffer) {
       this.logger.trace(`Data not found for contract ${contractAddress.toString()} and slot ${slot.toString()}`);
       return null;
@@ -57,8 +165,9 @@ export class CapsuleStore {
    * @param contractAddress - The contract address under which the data is scoped.
    * @param slot - The slot in the database to delete.
    */
-  async deleteCapsule(contractAddress: AztecAddress, slot: Fr): Promise<void> {
-    await this.#capsules.delete(dbSlotToKey(contractAddress, slot));
+  deleteCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string) {
+    // When we commit this, we will interpret null as a deletion, so we'll propagate the delete to the KV store
+    this.#deleteOnStage(jobId, dbSlotToKey(contractAddress, slot));
   }
 
   /**
@@ -72,13 +181,22 @@ export class CapsuleStore {
    * @param dstSlot - The first slot to copy to.
    * @param numEntries - The number of entries to copy.
    */
-  copyCapsule(contractAddress: AztecAddress, srcSlot: Fr, dstSlot: Fr, numEntries: number): Promise<void> {
+  copyCapsule(
+    contractAddress: AztecAddress,
+    srcSlot: Fr,
+    dstSlot: Fr,
+    numEntries: number,
+    jobId: string,
+  ): Promise<void> {
+    // This transactional context gives us "copy atomicity":
+    // there shouldn't be concurrent writes to what's being copied here.
+    // Equally important: this in practice is expected to perform thousands of DB operations
+    // and not using a transaction here would heavily impact performance.
     return this.#store.transactionAsync(async () => {
       // In order to support overlapping source and destination regions, we need to check the relative positions of source
       // and destination. If destination is ahead of source, then by the time we overwrite source elements using forward
       // indexes we'll have already read those. On the contrary, if source is ahead of destination we need to use backward
       // indexes to avoid reading elements that've been overwritten.
-
       const indexes = Array.from(Array(numEntries).keys());
       if (srcSlot.lt(dstSlot)) {
         indexes.reverse();
@@ -88,12 +206,12 @@ export class CapsuleStore {
         const currentSrcSlot = dbSlotToKey(contractAddress, srcSlot.add(new Fr(i)));
         const currentDstSlot = dbSlotToKey(contractAddress, dstSlot.add(new Fr(i)));
 
-        const toCopy = await this.#capsules.getAsync(currentSrcSlot);
+        const toCopy = await this.#getFromStage(jobId, currentSrcSlot);
         if (!toCopy) {
           throw new Error(`Attempted to copy empty slot ${currentSrcSlot} for contract ${contractAddress.toString()}`);
         }
 
-        await this.#capsules.set(currentDstSlot, toCopy);
+        this.#setOnStage(jobId, currentDstSlot, toCopy);
       }
     });
   }
@@ -106,35 +224,45 @@ export class CapsuleStore {
    * @param baseSlot - The slot where the array length is stored
    * @param content - Array of capsule data to append
    */
-  appendToCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]): Promise<void> {
+  appendToCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][], jobId: string): Promise<void> {
+    // We wrap this in a transaction to serialize concurrent calls from Promise.all.
+    // Without this, concurrent appends to the same array could race: both read length=0,
+    // both write at the same slots, one overwrites the other.
+    // Equally important: this in practice is expected to perform thousands of DB operations
+    // and not using a transaction here would heavily impact performance.
     return this.#store.transactionAsync(async () => {
       // Load current length, defaulting to 0 if not found
-      const lengthData = await this.loadCapsule(contractAddress, baseSlot);
+      const lengthData = await this.loadCapsule(contractAddress, baseSlot, jobId);
       const currentLength = lengthData ? lengthData[0].toNumber() : 0;
 
       // Store each capsule at consecutive slots after baseSlot + 1 + currentLength
       for (let i = 0; i < content.length; i++) {
         const nextSlot = arraySlot(baseSlot, currentLength + i);
-        await this.storeCapsule(contractAddress, nextSlot, content[i]);
+        this.storeCapsule(contractAddress, nextSlot, content[i], jobId);
       }
 
       // Update length to include all new capsules
       const newLength = currentLength + content.length;
-      await this.storeCapsule(contractAddress, baseSlot, [new Fr(newLength)]);
+      this.storeCapsule(contractAddress, baseSlot, [new Fr(newLength)], jobId);
     });
   }
 
-  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr): Promise<Fr[][]> {
+  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, jobId: string): Promise<Fr[][]> {
+    // I'm leaving this transactional context here though because I'm assuming this
+    // gives us "read array atomicity": there shouldn't be concurrent writes to what's being copied
+    // here.
+    // This is one point we should revisit in the future if we want to relax the concurrency
+    // of jobs: different calls running concurrently on the same contract may cause trouble.
     return this.#store.transactionAsync(async () => {
       // Load length, defaulting to 0 if not found
-      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const maybeLength = await this.loadCapsule(contractAddress, baseSlot, jobId);
       const length = maybeLength ? maybeLength[0].toBigInt() : 0n;
 
       const values: Fr[][] = [];
 
       // Read each capsule at consecutive slots after baseSlot
       for (let i = 0; i < length; i++) {
-        const currentValue = await this.loadCapsule(contractAddress, arraySlot(baseSlot, i));
+        const currentValue = await this.loadCapsule(contractAddress, arraySlot(baseSlot, i), jobId);
         if (currentValue == undefined) {
           throw new Error(
             `Expected non-empty value at capsule array in base slot ${baseSlot} at index ${i} for contract ${contractAddress}`,
@@ -148,23 +276,31 @@ export class CapsuleStore {
     });
   }
 
-  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]) {
+  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][], jobId: string) {
+    // This transactional context in theory isn't so critical now because we aren't
+    // writing to DB so if there's exceptions midway and it blows up, no visible impact
+    // to persistent storage will happen.
+    // I'm leaving this transactional context here though because I'm assuming this
+    // gives us "write array atomicity": there shouldn't be concurrent writes to what's being copied
+    // here.
+    // This is one point we should revisit in the future if we want to relax the concurrency
+    // of jobs: different calls running concurrently on the same contract may cause trouble.
     return this.#store.transactionAsync(async () => {
       // Load current length, defaulting to 0 if not found
-      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const maybeLength = await this.loadCapsule(contractAddress, baseSlot, jobId);
       const originalLength = maybeLength ? maybeLength[0].toNumber() : 0;
 
       // Set the new length
-      await this.storeCapsule(contractAddress, baseSlot, [new Fr(content.length)]);
+      this.storeCapsule(contractAddress, baseSlot, [new Fr(content.length)], jobId);
 
       // Store the new content, possibly overwriting existing values
       for (let i = 0; i < content.length; i++) {
-        await this.storeCapsule(contractAddress, arraySlot(baseSlot, i), content[i]);
+        this.storeCapsule(contractAddress, arraySlot(baseSlot, i), content[i], jobId);
       }
 
       // Clear any stragglers
       for (let i = content.length; i < originalLength; i++) {
-        await this.deleteCapsule(contractAddress, arraySlot(baseSlot, i));
+        this.deleteCapsule(contractAddress, arraySlot(baseSlot, i), jobId);
       }
     });
   }

package/src/storage/note_store/note_store.ts

@@ -124,10 +124,12 @@ export class NoteStore {
    * specified block number. It restores any notes that were nullified after the given block
    * and deletes any active notes created after that block.
    *
+   * IMPORTANT: This method must be called within a transaction to ensure atomicity.
+   *
    * @param blockNumber - The new chain tip after a reorg
    * @param synchedBlockNumber - The block number up to which PXE managed to sync before the reorg happened.
    */
-  public async rollbackNotesAndNullifiers(blockNumber: number, synchedBlockNumber: number): Promise<void> {
+  public async rollback(blockNumber: number, synchedBlockNumber: number): Promise<void> {
     await this.#rewindNullifiersAfterBlock(blockNumber, synchedBlockNumber);
     await this.#deleteActiveNotesAfterBlock(blockNumber);
   }
@@ -140,24 +142,22 @@ export class NoteStore {
    *
    * @param blockNumber - Notes created after this block number will be deleted
    */
-  #deleteActiveNotesAfterBlock(blockNumber: number): Promise<void> {
-    return this.#store.transactionAsync(async () => {
-      const notes = await toArray(this.#notes.valuesAsync());
-      for (const note of notes) {
-        const noteDao = NoteDao.fromBuffer(note);
-        if (noteDao.l2BlockNumber > blockNumber) {
-          const noteIndex = toBufferBE(noteDao.index, 32).toString('hex');
-          await this.#notes.delete(noteIndex);
-          await this.#notesToScope.delete(noteIndex);
-          await this.#nullifierToNoteId.delete(noteDao.siloedNullifier.toString());
-          const scopes = await toArray(this.#scopes.keysAsync());
-          for (const scope of scopes) {
-            await this.#notesByContractAndScope.get(scope)!.deleteValue(noteDao.contractAddress.toString(), noteIndex);
-            await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(noteDao.storageSlot.toString(), noteIndex);
-          }
+  async #deleteActiveNotesAfterBlock(blockNumber: number): Promise<void> {
+    const notes = await toArray(this.#notes.valuesAsync());
+    for (const note of notes) {
+      const noteDao = NoteDao.fromBuffer(note);
+      if (noteDao.l2BlockNumber > blockNumber) {
+        const noteIndex = toBufferBE(noteDao.index, 32).toString('hex');
+        await this.#notes.delete(noteIndex);
+        await this.#notesToScope.delete(noteIndex);
+        await this.#nullifierToNoteId.delete(noteDao.siloedNullifier.toString());
+        const scopes = await toArray(this.#scopes.keysAsync());
+        for (const scope of scopes) {
+          await this.#notesByContractAndScope.get(scope)!.deleteValue(noteDao.contractAddress.toString(), noteIndex);
+          await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(noteDao.storageSlot.toString(), noteIndex);
         }
       }
-    });
+    }
   }
 
   /**
@@ -171,50 +171,52 @@ export class NoteStore {
    * @param synchedBlockNumber - Upper bound for the block range to process
    */
   async #rewindNullifiersAfterBlock(blockNumber: number, synchedBlockNumber: number): Promise<void> {
-    await this.#store.transactionAsync(async () => {
-      const nullifiersToUndo: string[] = [];
-      const currentBlockNumber = blockNumber + 1;
-      for (let i = currentBlockNumber; i <= synchedBlockNumber; i++) {
-        nullifiersToUndo.push(...(await toArray(this.#nullifiersByBlockNumber.getValuesAsync(i))));
-      }
-      const notesIndexesToReinsert = await Promise.all(
-        nullifiersToUndo.map(nullifier => this.#nullifiedNotesByNullifier.getAsync(nullifier)),
-      );
-      const notNullNoteIndexes = notesIndexesToReinsert.filter(noteIndex => noteIndex != undefined);
-      const nullifiedNoteBuffers = await Promise.all(
-        notNullNoteIndexes.map(noteIndex => this.#nullifiedNotes.getAsync(noteIndex!)),
-      );
-      const noteDaos = nullifiedNoteBuffers
-        .filter(buffer => buffer != undefined)
-        .map(buffer => NoteDao.fromBuffer(buffer!));
-
-      for (const dao of noteDaos) {
-        const noteIndex = toBufferBE(dao.index, 32).toString('hex');
-        await this.#notes.set(noteIndex, dao.toBuffer());
-        await this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex);
+    const nullifiersToUndo: string[] = [];
+    const currentBlockNumber = blockNumber + 1;
+    for (let i = currentBlockNumber; i <= synchedBlockNumber; i++) {
+      nullifiersToUndo.push(...(await toArray(this.#nullifiersByBlockNumber.getValuesAsync(i))));
+    }
+    const notesIndexesToReinsert = await Promise.all(
+      nullifiersToUndo.map(nullifier => this.#nullifiedNotesByNullifier.getAsync(nullifier)),
+    );
+    const notNullNoteIndexes = notesIndexesToReinsert.filter(noteIndex => noteIndex != undefined);
+    const nullifiedNoteBuffers = await Promise.all(
+      notNullNoteIndexes.map(noteIndex => this.#nullifiedNotes.getAsync(noteIndex!)),
+    );
+    const noteDaos = nullifiedNoteBuffers
+      .filter(buffer => buffer != undefined)
+      .map(buffer => NoteDao.fromBuffer(buffer!));
 
-        const scopes = await toArray(this.#nullifiedNotesToScope.getValuesAsync(noteIndex));
+    for (const dao of noteDaos) {
+      const noteIndex = toBufferBE(dao.index, 32).toString('hex');
 
-        if (scopes.length === 0) {
-          // We should never run into this error because notes always have a scope assigned to them - either on initial
-          // insertion via `addNotes` or when removing their nullifiers.
-          throw new Error(`No scopes found for nullified note with index ${noteIndex}`);
-        }
+      const scopes = await toArray(this.#nullifiedNotesToScope.getValuesAsync(noteIndex));
 
-        for (const scope of scopes) {
-          await this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex);
-          await this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex);
-          await this.#notesToScope.set(noteIndex, scope);
-        }
+      if (scopes.length === 0) {
+        // We should never run into this error because notes always have a scope assigned to them - either on initial
+        // insertion via `addNotes` or when removing their nullifiers.
+        throw new Error(`No scopes found for nullified note with index ${noteIndex}`);
+      }
 
-        await this.#nullifiedNotes.delete(noteIndex);
-        await this.#nullifiedNotesToScope.delete(noteIndex);
-        await this.#nullifiersByBlockNumber.deleteValue(dao.l2BlockNumber, dao.siloedNullifier.toString());
-        await this.#nullifiedNotesByContract.deleteValue(dao.contractAddress.toString(), noteIndex);
-        await this.#nullifiedNotesByStorageSlot.deleteValue(dao.storageSlot.toString(), noteIndex);
-        await this.#nullifiedNotesByNullifier.delete(dao.siloedNullifier.toString());
+      for (const scope of scopes) {
+        await Promise.all([
+          this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex),
+          this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex),
+          this.#notesToScope.set(noteIndex, scope),
+        ]);
       }
-    });
+
+      await Promise.all([
+        this.#notes.set(noteIndex, dao.toBuffer()),
+        this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex),
+        this.#nullifiedNotes.delete(noteIndex),
+        this.#nullifiedNotesToScope.delete(noteIndex),
+        this.#nullifiersByBlockNumber.deleteValue(dao.l2BlockNumber, dao.siloedNullifier.toString()),
+        this.#nullifiedNotesByContract.deleteValue(dao.contractAddress.toString(), noteIndex),
+        this.#nullifiedNotesByStorageSlot.deleteValue(dao.storageSlot.toString(), noteIndex),
+        this.#nullifiedNotesByNullifier.delete(dao.siloedNullifier.toString()),
+      ]);
+    }
   }
 
   /**

package/src/storage/private_event_store/private_event_store.ts

@@ -174,43 +174,43 @@ export class PrivateEventStore {
   /**
    * Rolls back private events that were stored after a given `blockNumber` and up to `synchedBlockNumber` (the block
    * number up to which PXE managed to sync before the reorg happened).
+   *
+   * IMPORTANT: This method must be called within a transaction to ensure atomicity.
    */
-  public async rollbackEventsAfterBlock(blockNumber: number, synchedBlockNumber: number): Promise<void> {
-    await this.#store.transactionAsync(async () => {
-      let removedCount = 0;
-
-      for (let block = blockNumber + 1; block <= synchedBlockNumber; block++) {
-        const indices = await this.#eventsByBlockNumber.getAsync(block);
-        if (indices) {
-          await this.#eventsByBlockNumber.delete(block);
-
-          for (const eventCommitmentIndex of indices) {
-            const entry = await this.#eventLogs.getAsync(eventCommitmentIndex);
-            if (!entry) {
-              throw new Error(`Event log not found for eventCommitmentIndex ${eventCommitmentIndex}`);
-            }
-
-            await this.#eventLogs.delete(eventCommitmentIndex);
-            await this.#seenLogs.delete(eventCommitmentIndex);
-
-            // Update #eventsByContractScopeSelector using the stored lookupKey
-            const existingIndices = await this.#eventsByContractScopeSelector.getAsync(entry.lookupKey);
-            if (!existingIndices || existingIndices.length === 0) {
-              throw new Error(`No indices found in #eventsByContractScopeSelector for key ${entry.lookupKey}`);
-            }
-            const filteredIndices = existingIndices.filter(idx => idx !== eventCommitmentIndex);
-            if (filteredIndices.length === 0) {
-              await this.#eventsByContractScopeSelector.delete(entry.lookupKey);
-            } else {
-              await this.#eventsByContractScopeSelector.set(entry.lookupKey, filteredIndices);
-            }
-
-            removedCount++;
+  public async rollback(blockNumber: number, synchedBlockNumber: number): Promise<void> {
+    let removedCount = 0;
+
+    for (let block = blockNumber + 1; block <= synchedBlockNumber; block++) {
+      const indices = await this.#eventsByBlockNumber.getAsync(block);
+      if (indices) {
+        await this.#eventsByBlockNumber.delete(block);
+
+        for (const eventCommitmentIndex of indices) {
+          const entry = await this.#eventLogs.getAsync(eventCommitmentIndex);
+          if (!entry) {
+            throw new Error(`Event log not found for eventCommitmentIndex ${eventCommitmentIndex}`);
           }
+
+          await this.#eventLogs.delete(eventCommitmentIndex);
+          await this.#seenLogs.delete(eventCommitmentIndex);
+
+          // Update #eventsByContractScopeSelector using the stored lookupKey
+          const existingIndices = await this.#eventsByContractScopeSelector.getAsync(entry.lookupKey);
+          if (!existingIndices || existingIndices.length === 0) {
+            throw new Error(`No indices found in #eventsByContractScopeSelector for key ${entry.lookupKey}`);
+          }
+          const filteredIndices = existingIndices.filter(idx => idx !== eventCommitmentIndex);
+          if (filteredIndices.length === 0) {
+            await this.#eventsByContractScopeSelector.delete(entry.lookupKey);
+          } else {
+            await this.#eventsByContractScopeSelector.set(entry.lookupKey, filteredIndices);
+          }
+
+          removedCount++;
         }
       }
+    }
 
-      this.logger.verbose(`Rolled back ${removedCount} private events after block ${blockNumber}`);
-    });
+    this.logger.verbose(`Rolled back ${removedCount} private events after block ${blockNumber}`);
   }
 }

package/src/tagging/recipient_sync/load_private_logs_for_sender_recipient_pair.ts

@@ -68,7 +68,10 @@ export async function loadPrivateLogsForSenderRecipientPair(
       throw new Error('Node failed to return latest block header when syncing logs');
     }
 
-    [finalizedBlockNumber, currentTimestamp] = [l2Tips.finalized.number, latestBlockHeader.globalVariables.timestamp];
+    [finalizedBlockNumber, currentTimestamp] = [
+      l2Tips.finalized.block.number,
+      latestBlockHeader.globalVariables.timestamp,
+    ];
   }
 
   let start: number, end: number;

package/src/tagging/sender_sync/utils/get_status_change_of_pending.ts

@@ -10,7 +10,7 @@ export async function getStatusChangeOfPending(
   aztecNode: AztecNode,
 ): Promise<{ txHashesToFinalize: TxHash[]; txHashesToDrop: TxHash[] }> {
   // Get receipts for all pending tx hashes and the finalized block number.
-  const [receipts, { finalized }] = await Promise.all([
+  const [receipts, tips] = await Promise.all([
     Promise.all(pending.map(pendingTxHash => aztecNode.getTxReceipt(pendingTxHash))),
     aztecNode.getL2Tips(),
   ]);
@@ -22,7 +22,11 @@ export async function getStatusChangeOfPending(
     const receipt = receipts[i];
     const txHash = pending[i];
 
-    if (receipt.status === TxStatus.SUCCESS && receipt.blockNumber && receipt.blockNumber <= finalized.number) {
+    if (
+      receipt.status === TxStatus.SUCCESS &&
+      receipt.blockNumber &&
+      receipt.blockNumber <= tips.finalized.block.number
+    ) {
       // Tx has been included in a block and the corresponding block is finalized --> we mark the indexes as
       // finalized.
       txHashesToFinalize.push(txHash);