feat: Adding support for siloing notes in pxe database

yarn-project/circuit-types/src/interfaces/pxe.ts

@@ -250,8 +250,9 @@ export interface PXE {
    * Adds a note to the database.
    * @throws If the note hash of the note doesn't exist in the tree.
    * @param note - The note to add.
+   * @param scope - The scope to add the note under. Currently optional.
    */
-  addNote(note: ExtendedNote): Promise<void>;
+  addNote(note: ExtendedNote, scope?: AztecAddress): Promise<void>;
 
   /**
    * Adds a nullified note to the database.

yarn-project/circuit-types/src/notes/incoming_notes_filter.ts

@@ -20,4 +20,6 @@ export type IncomingNotesFilter = {
   status?: NoteStatus;
   /** The siloed nullifier for the note. */
   siloedNullifier?: Fr;
+  /** The scopes in which to get incoming notes from. This defaults to all scopes. */
+  scopes?: AztecAddress[];
 };

yarn-project/pxe/src/database/kv_pxe_database.ts

@@ -14,6 +14,7 @@ import {
   type AztecKVStore,
   type AztecMap,
   type AztecMultiMap,
+  type AztecSet,
   type AztecSingleton,
 } from '@aztec/kv-store';
 import { contractArtifactFromBuffer, contractArtifactToBuffer } from '@aztec/types/abi';
@@ -36,10 +37,7 @@ export class KVPxeDatabase implements PxeDatabase {
   #notes: AztecMap<string, Buffer>;
   #nullifiedNotes: AztecMap<string, Buffer>;
   #nullifierToNoteId: AztecMap<string, string>;
-  #notesByContract: AztecMultiMap<string, string>;
-  #notesByStorageSlot: AztecMultiMap<string, string>;
-  #notesByTxHash: AztecMultiMap<string, string>;
-  #notesByIvpkM: AztecMultiMap<string, string>;
+
   #nullifiedNotesByContract: AztecMultiMap<string, string>;
   #nullifiedNotesByStorageSlot: AztecMultiMap<string, string>;
   #nullifiedNotesByTxHash: AztecMultiMap<string, string>;
@@ -57,6 +55,12 @@ export class KVPxeDatabase implements PxeDatabase {
   #outgoingNotesByTxHash: AztecMultiMap<string, string>;
   #outgoingNotesByOvpkM: AztecMultiMap<string, string>;
 
+  #scopes: AztecSet<string>;
+  #notesByContractAndScope: Map<string, AztecMultiMap<string, string>>;
+  #notesByStorageSlotAndScope: Map<string, AztecMultiMap<string, string>>;
+  #notesByTxHashAndScope: Map<string, AztecMultiMap<string, string>>;
+  #notesByIvpkMAndScope: Map<string, AztecMultiMap<string, string>>;
+
   constructor(private db: AztecKVStore) {
     this.#db = db;
 
@@ -76,11 +80,6 @@ export class KVPxeDatabase implements PxeDatabase {
     this.#nullifiedNotes = db.openMap('nullified_notes');
     this.#nullifierToNoteId = db.openMap('nullifier_to_note');
 
-    this.#notesByContract = db.openMultiMap('notes_by_contract');
-    this.#notesByStorageSlot = db.openMultiMap('notes_by_storage_slot');
-    this.#notesByTxHash = db.openMultiMap('notes_by_tx_hash');
-    this.#notesByIvpkM = db.openMultiMap('notes_by_ivpk_m');
-
     this.#nullifiedNotesByContract = db.openMultiMap('nullified_notes_by_contract');
     this.#nullifiedNotesByStorageSlot = db.openMultiMap('nullified_notes_by_storage_slot');
     this.#nullifiedNotesByTxHash = db.openMultiMap('nullified_notes_by_tx_hash');
@@ -94,6 +93,19 @@ export class KVPxeDatabase implements PxeDatabase {
     this.#outgoingNotesByStorageSlot = db.openMultiMap('outgoing_notes_by_storage_slot');
     this.#outgoingNotesByTxHash = db.openMultiMap('outgoing_notes_by_tx_hash');
     this.#outgoingNotesByOvpkM = db.openMultiMap('outgoing_notes_by_ovpk_m');
+
+    this.#scopes = db.openSet('scopes');
+    this.#notesByContractAndScope = new Map<string, AztecMultiMap<string, string>>();
+    this.#notesByStorageSlotAndScope = new Map<string, AztecMultiMap<string, string>>();
+    this.#notesByTxHashAndScope = new Map<string, AztecMultiMap<string, string>>();
+    this.#notesByIvpkMAndScope = new Map<string, AztecMultiMap<string, string>>();
+
+    for (const scope of this.#scopes.entries()) {
+      this.#notesByContractAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_contract`));
+      this.#notesByStorageSlotAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_storage_slot`));
+      this.#notesByTxHashAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_tx_hash`));
+      this.#notesByIvpkMAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_ivpk_m`));
+    }
   }
 
   public async getContract(
@@ -154,11 +166,19 @@ export class KVPxeDatabase implements PxeDatabase {
     return val?.map(b => Fr.fromBuffer(b));
   }
 
-  async addNote(note: IncomingNoteDao): Promise<void> {
-    await this.addNotes([note], []);
+  async addNote(note: IncomingNoteDao, scope?: AztecAddress): Promise<void> {
+    await this.addNotes([note], [], scope);
   }
 
-  addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise<void> {
+  async addNotes(
+    incomingNotes: IncomingNoteDao[],
+    outgoingNotes: OutgoingNoteDao[],
+    scope: AztecAddress = AztecAddress.ZERO,
+  ): Promise<void> {
+    if (!this.#scopes.has(scope.toString())) {
+      await this.#addScope(scope);
+    }
+
     return this.db.transaction(() => {
       for (const dao of incomingNotes) {
         // store notes by their index in the notes hash tree
@@ -168,10 +188,11 @@ export class KVPxeDatabase implements PxeDatabase {
         const noteIndex = toBufferBE(dao.index, 32).toString('hex');
         void this.#notes.set(noteIndex, dao.toBuffer());
         void this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex);
-        void this.#notesByContract.set(dao.contractAddress.toString(), noteIndex);
-        void this.#notesByStorageSlot.set(dao.storageSlot.toString(), noteIndex);
-        void this.#notesByTxHash.set(dao.txHash.toString(), noteIndex);
-        void this.#notesByIvpkM.set(dao.ivpkM.toString(), noteIndex);
+
+        void this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex);
+        void this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex);
+        void this.#notesByTxHashAndScope.get(scope.toString())!.set(dao.txHash.toString(), noteIndex);
+        void this.#notesByIvpkMAndScope.get(scope.toString())!.set(dao.ivpkM.toString(), noteIndex);
       }
 
       for (const dao of outgoingNotes) {
@@ -244,16 +265,31 @@ export class KVPxeDatabase implements PxeDatabase {
 
     const candidateNoteSources = [];
 
+    filter.scopes ??= [...this.#scopes.entries()].map(addressString => AztecAddress.fromString(addressString));
+
+    const activeNoteIdsPerScope: IterableIterator<string>[] = [];
+
+    for (const scope of new Set(filter.scopes)) {
+      const formattedScopeString = scope.toString();
+      if (!this.#scopes.has(formattedScopeString)) {
+        throw new Error('Trying to get incoming notes of an scope that is not in the PXE database');
+      }
+
+      activeNoteIdsPerScope.push(
+        publicKey
+          ? this.#notesByIvpkMAndScope.get(formattedScopeString)!.getValues(publicKey.toString())
+          : filter.txHash
+          ? this.#notesByTxHashAndScope.get(formattedScopeString)!.getValues(filter.txHash.toString())
+          : filter.contractAddress
+          ? this.#notesByContractAndScope.get(formattedScopeString)!.getValues(filter.contractAddress.toString())
+          : filter.storageSlot
+          ? this.#notesByStorageSlotAndScope.get(formattedScopeString)!.getValues(filter.storageSlot.toString())
+          : this.#notesByIvpkMAndScope.get(formattedScopeString)!.values(),
+      );
+    }
+
     candidateNoteSources.push({
-      ids: publicKey
-        ? this.#notesByIvpkM.getValues(publicKey.toString())
-        : filter.txHash
-        ? this.#notesByTxHash.getValues(filter.txHash.toString())
-        : filter.contractAddress
-        ? this.#notesByContract.getValues(filter.contractAddress.toString())
-        : filter.storageSlot
-        ? this.#notesByStorageSlot.getValues(filter.storageSlot.toString())
-        : this.#notes.keys(),
+      ids: new Set(activeNoteIdsPerScope.flatMap(iterableIterator => [...iterableIterator])),
       notes: this.#notes,
     });
 
@@ -358,7 +394,7 @@ export class KVPxeDatabase implements PxeDatabase {
     return Promise.resolve(notes);
   }
 
-  removeNullifiedNotes(nullifiers: Fr[], account: PublicKey): Promise<IncomingNoteDao[]> {
+  removeNullifiedNotes(nullifiers: Fr[], accountIvpkM: PublicKey): Promise<IncomingNoteDao[]> {
     if (nullifiers.length === 0) {
       return Promise.resolve([]);
     }
@@ -380,18 +416,21 @@ export class KVPxeDatabase implements PxeDatabase {
         }
 
         const note = IncomingNoteDao.fromBuffer(noteBuffer);
-        if (!note.ivpkM.equals(account)) {
+        if (!note.ivpkM.equals(accountIvpkM)) {
           // tried to nullify someone else's note
           continue;
         }
 
         nullifiedNotes.push(note);
 
         void this.#notes.delete(noteIndex);
-        void this.#notesByIvpkM.deleteValue(account.toString(), noteIndex);
-        void this.#notesByTxHash.deleteValue(note.txHash.toString(), noteIndex);
-        void this.#notesByContract.deleteValue(note.contractAddress.toString(), noteIndex);
-        void this.#notesByStorageSlot.deleteValue(note.storageSlot.toString(), noteIndex);
+
+        for (const scope in this.#scopes.entries()) {
+          void this.#notesByIvpkMAndScope.get(scope)!.deleteValue(accountIvpkM.toString(), noteIndex);
+          void this.#notesByTxHashAndScope.get(scope)!.deleteValue(note.txHash.toString(), noteIndex);
+          void this.#notesByContractAndScope.get(scope)!.deleteValue(note.contractAddress.toString(), noteIndex);
+          void this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(note.storageSlot.toString(), noteIndex);
+        }
 
         void this.#nullifiedNotes.set(noteIndex, note.toBuffer());
         void this.#nullifiedNotesByContract.set(note.contractAddress.toString(), noteIndex);
@@ -440,6 +479,22 @@ export class KVPxeDatabase implements PxeDatabase {
     return Header.fromBuffer(headerBuffer);
   }
 
+  async #addScope(scope: AztecAddress): Promise<boolean> {
+    const scopeString = scope.toString();
+
+    if (this.#scopes.has(scopeString)) {
+      return false;
+    }
+
+    await this.#scopes.add(scopeString);
+    this.#notesByContractAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_contract`));
+    this.#notesByStorageSlotAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_storage_slot`));
+    this.#notesByTxHashAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_tx_hash`));
+    this.#notesByIvpkMAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_ivpk_m`));
+
+    return true;
+  }
+
   addCompleteAddress(completeAddress: CompleteAddress): Promise<boolean> {
     return this.#db.transaction(() => {
       const addressString = completeAddress.address.toString();

yarn-project/pxe/src/database/pxe_database.ts

@@ -62,8 +62,10 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD
   /**
    * Adds a note to DB.
    * @param note - The note to add.
+   * @param scope - The scope to add the note under. Currently optional.
+   * @remark - Will create a database for the scope if it does not already exist.
    */
-  addNote(note: IncomingNoteDao): Promise<void>;
+  addNote(note: IncomingNoteDao, scope?: AztecAddress): Promise<void>;
 
   /**
    * Adds a nullified note to DB.
@@ -78,8 +80,10 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD
    *
    * @param incomingNotes - An array of notes which were decrypted as incoming.
    * @param outgoingNotes - An array of notes which were decrypted as outgoing.
+   * @param scope - The scope to add the notes under. Currently optional.
+   * @remark - Will create a database for the scope if it does not already exist.
    */
-  addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise<void>;
+  addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[], scope?: AztecAddress): Promise<void>;
 
   /**
    * Add notes to the database that are intended for us, but we don't yet have the contract.

yarn-project/pxe/src/database/pxe_database_test_suite.ts

@@ -130,14 +130,19 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
 
       it.each(filteringTests)('stores notes in bulk and retrieves notes', async (getFilter, getExpected) => {
         await database.addNotes(notes, []);
-        await expect(database.getIncomingNotes(getFilter())).resolves.toEqual(getExpected());
+        const returnedNotes = await database.getIncomingNotes(getFilter());
+
+        expect(returnedNotes.sort()).toEqual(getExpected().sort());
       });
 
       it.each(filteringTests)('stores notes one by one and retrieves notes', async (getFilter, getExpected) => {
         for (const note of notes) {
           await database.addNote(note);
         }
-        await expect(database.getIncomingNotes(getFilter())).resolves.toEqual(getExpected());
+
+        const returnedNotes = await database.getIncomingNotes(getFilter());
+
+        expect(returnedNotes.sort()).toEqual(getExpected().sort());
       });
 
       it.each(filteringTests)('retrieves nullified notes', async (getFilter, getExpected) => {
@@ -196,6 +201,68 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
         // inserted combining active and nullified results.
         expect(result.sort()).toEqual([...notes].sort());
       });
+
+      it('stores notes one by one and retrieves notes with siloed account', async () => {
+        for (const note of notes.slice(0, 5)) {
+          await database.addNote(note, owners[0].address);
+        }
+
+        for (const note of notes.slice(5)) {
+          await database.addNote(note, owners[1].address);
+        }
+
+        const owner0IncomingNotes = await database.getIncomingNotes({
+          scopes: [owners[0].address],
+        });
+
+        expect(owner0IncomingNotes.sort()).toEqual(notes.slice(0, 5).sort());
+
+        const owner1IncomingNotes = await database.getIncomingNotes({
+          scopes: [owners[1].address],
+        });
+
+        expect(owner1IncomingNotes.sort()).toEqual(notes.slice(5).sort());
+
+        const bothOwnerIncomingNotes = await database.getIncomingNotes({
+          scopes: [owners[0].address, owners[1].address],
+        });
+
+        expect(bothOwnerIncomingNotes.sort()).toEqual(notes.sort());
+      });
+
+      it('a nullified note removes notes from all accounts in the pxe', async () => {
+        await database.addNote(notes[0], owners[0].address);
+        await database.addNote(notes[0], owners[1].address);
+
+        await expect(
+          database.getIncomingNotes({
+            scopes: [owners[0].address],
+          }),
+        ).resolves.toEqual([notes[0]]);
+        await expect(
+          database.getIncomingNotes({
+            scopes: [owners[1].address],
+          }),
+        ).resolves.toEqual([notes[0]]);
+
+        await expect(
+          database.removeNullifiedNotes(
+            [notes[0].siloedNullifier],
+            owners[0].publicKeys.masterIncomingViewingPublicKey,
+          ),
+        ).resolves.toEqual([notes[0]]);
+
+        await expect(
+          database.getIncomingNotes({
+            scopes: [owners[0].address],
+          }),
+        ).resolves.toEqual([]);
+        await expect(
+          database.getIncomingNotes({
+            scopes: [owners[1].address],
+          }),
+        ).resolves.toEqual([]);
+      });
     });
 
     describe('outgoing notes', () => {

yarn-project/pxe/src/pxe_service/pxe_service.ts

@@ -338,7 +338,7 @@ export class PXEService implements PXE {
     return Promise.all(extendedNotes);
   }
 
-  public async addNote(note: ExtendedNote) {
+  public async addNote(note: ExtendedNote, scope?: AztecAddress) {
     const owner = await this.db.getCompleteAddress(note.owner);
     if (!owner) {
       throw new Error(`Unknown account: ${note.owner.toString()}`);
@@ -384,6 +384,7 @@ export class PXEService implements PXE {
           index,
           owner.publicKeys.masterIncomingViewingPublicKey,
         ),
+        scope,
       );
     }
   }

yarn-project/tsconfig.json

@@ -51,7 +51,7 @@
     { "path": "scripts/tsconfig.json" },
     { "path": "entrypoints/tsconfig.json" },
     { "path": "cli/tsconfig.json" },
-    { "path": "cli-wallet/tsconfig.json"}
+    { "path": "cli-wallet/tsconfig.json" }
   ],
   "files": ["./@types/jest/index.d.ts"],
   "exclude": ["node_modules", "**/node_modules", "**/.*/"]