@aztec/pxe 0.0.1-commit.24de95ac → 0.0.1-commit.3469e52

package/src/storage/{note_data_provider/note_data_provider.ts → note_store/note_store.ts}

@@ -1,47 +1,54 @@
-import { toBufferBE } from '@aztec/foundation/bigint-buffer';
-import type { Fr } from '@aztec/foundation/fields';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import { toArray } from '@aztec/foundation/iterable';
 import type { AztecAsyncKVStore, AztecAsyncMap, AztecAsyncMultiMap } from '@aztec/kv-store';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import type { InBlock } from '@aztec/stdlib/block';
+import type { DataInBlock } from '@aztec/stdlib/block';
 import { NoteStatus, type NotesFilter } from '@aztec/stdlib/note';
-
-import { NoteDao } from './note_dao.js';
+import { NoteDao } from '@aztec/stdlib/note';
 
 /**
- * NoteDataProvider manages the storage and retrieval of notes.
+ * NoteStore manages the storage and retrieval of notes.
  *
  * Notes can be active or nullified. This class processes new notes, nullifications,
  * and performs rollback handling in the case of a reorg.
  **/
-export class NoteDataProvider {
+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');
@@ -50,16 +57,16 @@ export class NoteDataProvider {
   }
 
   /**
-   * Creates and initializes a new NoteDataProvider instance.
+   * Creates and initializes a new NoteStore instance.
    *
-   * This factory method creates a NoteDataProvider and restores any existing
+   * This factory method creates a NoteStore and restores any existing
    * scope-specific indexes from the database.
    *
    * @param store - The key-value store to use for persistence
-   * @returns Promise resolving to a fully initialized NoteDataProvider instance
+   * @returns Promise resolving to a fully initialized NoteStore instance
    */
-  public static async create(store: AztecAsyncKVStore): Promise<NoteDataProvider> {
-    const pxeDB = new NoteDataProvider(store);
+  public static async create(store: AztecAsyncKVStore): Promise<NoteStore> {
+    const pxeDB = new NoteStore(store);
     for await (const scope of pxeDB.#scopes.keysAsync()) {
       pxeDB.#notesByContractAndScope.set(scope, store.openMultiMap(`${scope}:notes_by_contract`));
       pxeDB.#notesByStorageSlotAndScope.set(scope, store.openMultiMap(`${scope}:notes_by_storage_slot`));
@@ -93,9 +100,8 @@ export class NoteDataProvider {
   /**
    * 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
@@ -107,13 +113,12 @@ export class NoteDataProvider {
       }
 
       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);
       }
     });
   }
@@ -125,10 +130,12 @@ export class NoteDataProvider {
    * 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);
   }
@@ -141,24 +148,21 @@ export class NoteDataProvider {
    *
    * @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);
         }
       }
-    });
+    }
   }
 
   /**
@@ -172,48 +176,47 @@ export class NoteDataProvider {
    * @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!));
 
-        let scopes = (await toArray(this.#nullifiedNotesToScope.getValuesAsync(noteIndex))) ?? [];
+    for (const dao of noteDaos) {
+      const noteId = dao.siloedNullifier.toString();
 
-        if (scopes.length === 0) {
-          scopes = [dao.recipient.toString()];
-        }
+      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),
+      ]);
+    }
   }
 
   /**
@@ -223,8 +226,9 @@ export class NoteDataProvider {
    * parameters.
    *
    * @param filter - Filter criteria including contractAddress (required), and optional
-   *                 storageSlot, status, scopes and siloedNullifier.
-   * @returns Promise resolving to array of NoteDao objects matching the filter
+   *                 owner, storageSlot, status, scopes, and siloedNullifier.
+   * @returns Filtered and deduplicated notes (a note might be present in multiple scopes - we ensure it is only
+   * returned once if this is the case)
    * @throws If filtering by an empty scopes array. Scopes have to be set to undefined or to a non-empty array.
    */
   async getNotes(filter: NotesFilter): Promise<NoteDao[]> {
@@ -306,6 +310,10 @@ export class NoteDataProvider {
           continue;
         }
 
+        if (filter.owner && !note.owner.equals(filter.owner)) {
+          continue;
+        }
+
         if (filter.storageSlot && !note.storageSlot.equals(filter.storageSlot!)) {
           continue;
         }
@@ -318,7 +326,26 @@ export class NoteDataProvider {
       }
     }
 
-    return result;
+    // A note might be present in multiple scopes - we ensure it is only returned once
+    const deduplicated: NoteDao[] = [];
+    for (const note of result) {
+      if (!deduplicated.some(existing => existing.equals(note))) {
+        deduplicated.push(note);
+      }
+    }
+
+    // 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;
   }
 
   /**
@@ -332,7 +359,7 @@ export class NoteDataProvider {
    * @returns Promise resolving to array of nullified NoteDao objects
    * @throws Error if any nullifier is not found in the active notes
    */
-  applyNullifiers(nullifiers: InBlock<Fr>[]): Promise<NoteDao[]> {
+  applyNullifiers(nullifiers: DataInBlock<Fr>[]): Promise<NoteDao[]> {
     if (nullifiers.length === 0) {
       return Promise.resolve([]);
     }
@@ -342,50 +369,45 @@ export class NoteDataProvider {
 
       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(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.
+          throw new Error('Note scopes are missing in applyNullifiers');
         }
-        const noteScopes = (await toArray(this.#notesToScope.getValuesAsync(noteIndex))) ?? [];
+
         const note = NoteDao.fromBuffer(noteBuffer);
 
         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);
         }
 
-        if (noteScopes !== undefined) {
-          for (const scope of noteScopes) {
-            await this.#nullifiedNotesToScope.set(noteIndex, scope);
-          }
+        for (const scope of noteScopes) {
+          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;
     });