@aztec/pxe 4.0.0-nightly.20260112 → 4.0.0-nightly.20260114

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

@@ -1,4 +1,3 @@
-import { toBufferBE } from '@aztec/foundation/bigint-buffer';
 import type { Fr } from '@aztec/foundation/curves/bn254';
 import { toArray } from '@aztec/foundation/iterable';
 import type { AztecAsyncKVStore, AztecAsyncMap, AztecAsyncMultiMap } from '@aztec/kv-store';
@@ -15,32 +14,41 @@ import { NoteDao } from '@aztec/stdlib/note';
  **/
 export class NoteStore {
   #store: AztecAsyncKVStore;
+
+  // Note that we use the siloedNullifier as the note id in the store as it's guaranteed to be unique.
+
+  /** noteId (siloedNullifier) -> NoteDao (serialized) */
   #notes: AztecAsyncMap<string, Buffer>;
+  /** noteId (siloedNullifier) -> NoteDao (serialized) */
   #nullifiedNotes: AztecAsyncMap<string, Buffer>;
-  #nullifierToNoteId: AztecAsyncMap<string, string>;
+  /** blockNumber -> siloedNullifier */
   #nullifiersByBlockNumber: AztecAsyncMultiMap<number, string>;
 
+  /** noteId (siloedNullifier) -> scope */
   #nullifiedNotesToScope: AztecAsyncMultiMap<string, string>;
+  /** contractAddress -> noteId (siloedNullifier) */
   #nullifiedNotesByContract: AztecAsyncMultiMap<string, string>;
+  /** storageSlot -> noteId (siloedNullifier) */
   #nullifiedNotesByStorageSlot: AztecAsyncMultiMap<string, string>;
-  #nullifiedNotesByNullifier: AztecAsyncMap<string, string>;
 
+  /** scope (AztecAddress) -> true */
   #scopes: AztecAsyncMap<string, true>;
+  /** noteId (siloedNullifier) -> scope */
   #notesToScope: AztecAsyncMultiMap<string, string>;
+  /** scope -> MultiMap(contractAddress -> noteId) */
   #notesByContractAndScope: Map<string, AztecAsyncMultiMap<string, string>>;
+  /** scope -> MultiMap(storageSlot -> noteId) */
   #notesByStorageSlotAndScope: Map<string, AztecAsyncMultiMap<string, string>>;
 
   private constructor(store: AztecAsyncKVStore) {
     this.#store = store;
     this.#notes = store.openMap('notes');
     this.#nullifiedNotes = store.openMap('nullified_notes');
-    this.#nullifierToNoteId = store.openMap('nullifier_to_note');
     this.#nullifiersByBlockNumber = store.openMultiMap('nullifier_to_block_number');
 
     this.#nullifiedNotesToScope = store.openMultiMap('nullified_notes_to_scope');
     this.#nullifiedNotesByContract = store.openMultiMap('nullified_notes_by_contract');
     this.#nullifiedNotesByStorageSlot = store.openMultiMap('nullified_notes_by_storage_slot');
-    this.#nullifiedNotesByNullifier = store.openMap('nullified_notes_by_nullifier');
 
     this.#scopes = store.openMap('scopes');
     this.#notesToScope = store.openMultiMap('notes_to_scope');
@@ -92,9 +100,8 @@ export class NoteStore {
   /**
    * Adds multiple notes to the data provider under the specified scope.
    *
-   * Notes are stored using their index from the notes hash tree as the key, which provides
-   * uniqueness and maintains creation order. Each note is indexed by multiple criteria
-   * for efficient retrieval.
+   * Notes are stored using their siloedNullifier as the key, which provides uniqueness. Each note is indexed
+   * by multiple criteria for efficient retrieval.
    *
    * @param notes - Notes to store
    * @param scope - The scope (user/account) under which to store the notes
@@ -106,13 +113,12 @@ export class NoteStore {
       }
 
       for (const dao of notes) {
-        const noteIndex = toBufferBE(dao.index, 32).toString('hex');
-        await this.#notes.set(noteIndex, dao.toBuffer());
-        await this.#notesToScope.set(noteIndex, scope.toString());
-        await this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex);
+        const noteId = dao.siloedNullifier.toString();
+        await this.#notes.set(noteId, dao.toBuffer());
+        await this.#notesToScope.set(noteId, scope.toString());
 
-        await this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex);
-        await this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex);
+        await this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteId);
+        await this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteId);
       }
     });
   }
@@ -124,10 +130,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 +148,21 @@ 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 noteId = noteDao.siloedNullifier.toString();
+        await this.#notes.delete(noteId);
+        await this.#notesToScope.delete(noteId);
+        const scopes = await toArray(this.#scopes.keysAsync());
+        for (const scope of scopes) {
+          await this.#notesByContractAndScope.get(scope)!.deleteValue(noteDao.contractAddress.toString(), noteId);
+          await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(noteDao.storageSlot.toString(), noteId);
         }
       }
-    });
+    }
   }
 
   /**
@@ -171,50 +176,47 @@ 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 noteIdsToReinsert: string[] = [];
+    const currentBlockNumber = blockNumber + 1;
+    for (let i = currentBlockNumber; i <= synchedBlockNumber; i++) {
+      // noteId === siloedNullifier.toString(), so we can use nullifiers directly as noteIds
+      noteIdsToReinsert.push(...(await toArray(this.#nullifiersByBlockNumber.getValuesAsync(i))));
+    }
+    const nullifiedNoteBuffers = await Promise.all(
+      noteIdsToReinsert.map(noteId => this.#nullifiedNotes.getAsync(noteId)),
+    );
+    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 noteId = dao.siloedNullifier.toString();
 
-        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(noteId));
 
-        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 nullifier ${noteId}`);
+      }
 
-        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(), noteId),
+          this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteId),
+          this.#notesToScope.set(noteId, scope),
+        ]);
       }
-    });
+
+      await Promise.all([
+        this.#notes.set(noteId, dao.toBuffer()),
+        this.#nullifiedNotes.delete(noteId),
+        this.#nullifiedNotesToScope.delete(noteId),
+        this.#nullifiersByBlockNumber.deleteValue(dao.l2BlockNumber, dao.siloedNullifier.toString()),
+        this.#nullifiedNotesByContract.deleteValue(dao.contractAddress.toString(), noteId),
+        this.#nullifiedNotesByStorageSlot.deleteValue(dao.storageSlot.toString(), noteId),
+      ]);
+    }
   }
 
   /**
@@ -332,6 +334,17 @@ export class NoteStore {
       }
     }
 
+    // Sort by block number, then by tx index within block, then by note index within tx
+    deduplicated.sort((a, b) => {
+      if (a.l2BlockNumber !== b.l2BlockNumber) {
+        return a.l2BlockNumber - b.l2BlockNumber;
+      }
+      if (a.txIndexInBlock !== b.txIndexInBlock) {
+        return a.txIndexInBlock - b.txIndexInBlock;
+      }
+      return a.noteIndexInTx - b.noteIndexInTx;
+    });
+
     return deduplicated;
   }
 
@@ -356,25 +369,18 @@ export class NoteStore {
 
       for (const blockScopedNullifier of nullifiers) {
         const { data: nullifier, l2BlockNumber: blockNumber } = blockScopedNullifier;
-        const nullifierKey = nullifier.toString();
+        const noteId = nullifier.toString();
 
-        const noteIndex = await this.#nullifierToNoteId.getAsync(nullifierKey);
-        if (!noteIndex) {
-          // Check if already nullified?
-          const alreadyNullified = await this.#nullifiedNotesByNullifier.getAsync(nullifierKey);
-          if (alreadyNullified) {
+        const noteBuffer = await this.#notes.getAsync(noteId);
+        if (!noteBuffer) {
+          // Check if already nullified (noteId === siloedNullifier, so we can check #nullifiedNotes directly)
+          if (await this.#nullifiedNotes.hasAsync(noteId)) {
             throw new Error(`Nullifier already applied in applyNullifiers`);
           }
           throw new Error('Nullifier not found in applyNullifiers');
         }
 
-        const noteBuffer = noteIndex ? await this.#notes.getAsync(noteIndex) : undefined;
-
-        if (!noteBuffer) {
-          throw new Error('Note not found in applyNullifiers');
-        }
-
-        const noteScopes = await toArray(this.#notesToScope.getValuesAsync(noteIndex));
+        const noteScopes = await toArray(this.#notesToScope.getValuesAsync(noteId));
         if (noteScopes.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.
@@ -385,26 +391,23 @@ export class NoteStore {
 
         nullifiedNotes.push(note);
 
-        await this.#notes.delete(noteIndex);
-        await this.#notesToScope.delete(noteIndex);
+        await this.#notes.delete(noteId);
+        await this.#notesToScope.delete(noteId);
 
         const scopes = await toArray(this.#scopes.keysAsync());
 
         for (const scope of scopes) {
-          await this.#notesByContractAndScope.get(scope)!.deleteValue(note.contractAddress.toString(), noteIndex);
-          await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(note.storageSlot.toString(), noteIndex);
+          await this.#notesByContractAndScope.get(scope)!.deleteValue(note.contractAddress.toString(), noteId);
+          await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(note.storageSlot.toString(), noteId);
         }
 
         for (const scope of noteScopes) {
-          await this.#nullifiedNotesToScope.set(noteIndex, scope);
+          await this.#nullifiedNotesToScope.set(noteId, scope);
         }
-        await this.#nullifiedNotes.set(noteIndex, note.toBuffer());
-        await this.#nullifiersByBlockNumber.set(blockNumber, nullifier.toString());
-        await this.#nullifiedNotesByContract.set(note.contractAddress.toString(), noteIndex);
-        await this.#nullifiedNotesByStorageSlot.set(note.storageSlot.toString(), noteIndex);
-        await this.#nullifiedNotesByNullifier.set(nullifier.toString(), noteIndex);
-
-        await this.#nullifierToNoteId.delete(nullifier.toString());
+        await this.#nullifiedNotes.set(noteId, note.toBuffer());
+        await this.#nullifiersByBlockNumber.set(blockNumber, noteId);
+        await this.#nullifiedNotesByContract.set(note.contractAddress.toString(), noteId);
+        await this.#nullifiedNotesByStorageSlot.set(note.storageSlot.toString(), noteId);
       }
       return nullifiedNotes;
     });