fix: bug in PXE::getNotes

yarn-project/cli/src/utils/inspect.ts

@@ -3,7 +3,7 @@ import type { LogFn } from '@aztec/foundation/log';
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
 import { siloNullifier } from '@aztec/stdlib/hash';
 import type { PXE } from '@aztec/stdlib/interfaces/client';
-import { type ExtendedNote, NoteStatus } from '@aztec/stdlib/note';
+import type { ExtendedNote } from '@aztec/stdlib/note';
 import type { TxHash } from '@aztec/stdlib/tx';
 
 export async function inspectBlock(pxe: PXE, blockNumber: number, log: LogFn, opts: { showTxs?: boolean } = {}) {
@@ -40,11 +40,7 @@ export async function inspectTx(
   log: LogFn,
   opts: { includeBlockInfo?: boolean; artifactMap?: ArtifactMap } = {},
 ) {
-  const [receipt, effectsInBlock, getNotes] = await Promise.all([
-    pxe.getTxReceipt(txHash),
-    pxe.getTxEffect(txHash),
-    pxe.getNotes({ txHash, status: NoteStatus.ACTIVE_OR_NULLIFIED }),
-  ]);
+  const [receipt, effectsInBlock] = await Promise.all([pxe.getTxReceipt(txHash), pxe.getTxEffect(txHash)]);
   // Base tx data
   log(`Tx ${txHash.toString()}`);
   log(` Status: ${receipt.status} ${effectsInBlock ? `(${effectsInBlock.data.revertCode.getDescription()})` : ''}`);
@@ -88,12 +84,9 @@ export async function inspectTx(
   const notes = effects.noteHashes;
   if (notes.length > 0) {
     log(' Created notes:');
-    log(`  Total: ${notes.length}. Found: ${getNotes.length}.`);
-    if (getNotes.length) {
-      log('  Found notes:');
-      for (const note of getNotes) {
-        inspectNote(note, artifactMap, log);
-      }
+    log(`  Total: ${notes.length}`);
+    for (const note of notes) {
+      log(`  Note hash: ${note.toShortString()}`);
     }
   }
 
@@ -103,8 +96,10 @@ export async function inspectTx(
   if (nullifierCount > 0) {
     log(' Nullifiers:');
     for (const nullifier of effects.nullifiers) {
-      const [note] = await pxe.getNotes({ siloedNullifier: nullifier });
       const deployed = deployNullifiers[nullifier.toString()];
+      const note = deployed
+        ? (await pxe.getNotes({ siloedNullifier: nullifier, contractAddress: deployed }))[0]
+        : undefined;
       const initialized = initNullifiers[nullifier.toString()];
       const registered = classNullifiers[nullifier.toString()];
       if (nullifier.toBuffer().equals(txHash.toBuffer())) {

yarn-project/end-to-end/src/e2e_blacklist_token_contract/minting.test.ts

@@ -93,10 +93,8 @@ describe('e2e_blacklist_token_contract mint', () => {
         const receiptClaim = await asset.methods.redeem_shield(wallets[0].getAddress(), amount, secret).send().wait();
 
         tokenSim.mintPrivate(wallets[0].getAddress(), amount);
-        // Trigger a note sync
-        await asset.methods.sync_private_state().simulate();
         // 1 note should have been created containing `amount` of tokens
-        const visibleNotes = await pxe.getNotes({ txHash: receiptClaim.txHash });
+        const visibleNotes = await pxe.getNotes({ txHash: receiptClaim.txHash, contractAddress: asset.address });
         expect(visibleNotes.length).toBe(1);
         expect(visibleNotes[0].note.items[0].toBigInt()).toBe(amount);
       });

yarn-project/end-to-end/src/e2e_crowdfunding_and_claim.test.ts

@@ -161,8 +161,10 @@ describe('e2e_crowdfunding_and_claim', () => {
         .wait();
 
       // Get the notes emitted by the Crowdfunding contract and check that only 1 was emitted (the UintNote)
-      await crowdfundingContract.withWallet(donorWallets[0]).methods.sync_private_state().simulate();
-      const notes = await pxe.getNotes({ txHash: donateTxReceipt.txHash });
+      const notes = await pxe.getNotes({
+        txHash: donateTxReceipt.txHash,
+        contractAddress: crowdfundingContract.address,
+      });
       const filteredNotes = notes.filter(x => x.contractAddress.equals(crowdfundingContract.address));
       expect(filteredNotes!.length).toEqual(1);
 
@@ -225,8 +227,7 @@ describe('e2e_crowdfunding_and_claim', () => {
       .wait();
 
     // Get the notes emitted by the Crowdfunding contract and check that only 1 was emitted (the UintNote)
-    await crowdfundingContract.withWallet(unrelatedWallet).methods.sync_private_state().simulate();
-    const notes = await pxe.getNotes({ txHash: donateTxReceipt.txHash });
+    const notes = await pxe.getNotes({ contractAddress: crowdfundingContract.address, txHash: donateTxReceipt.txHash });
     const filtered = notes.filter(x => x.contractAddress.equals(crowdfundingContract.address));
     expect(filtered!.length).toEqual(1);
 
@@ -272,8 +273,7 @@ describe('e2e_crowdfunding_and_claim', () => {
         .call_create_note(arbitraryValue, owner, sender, arbitraryStorageSlot)
         .send()
         .wait();
-      await testContract.methods.sync_private_state().simulate();
-      const notes = await pxe.getNotes({ txHash: receipt.txHash });
+      const notes = await pxe.getNotes({ txHash: receipt.txHash, contractAddress: testContract.address });
       expect(notes.length).toEqual(1);
       note = processUniqueNote(notes[0]);
     }

yarn-project/end-to-end/src/e2e_pending_note_hashes_contract.test.ts

@@ -293,9 +293,7 @@ describe('e2e_pending_note_hashes_contract', () => {
     // Then emit another note log with the same counter as the one above, but with value 5
     const txReceipt = await deployedContract.methods.test_emit_bad_note_log(owner, sender).send().wait();
 
-    await deployedContract.methods.sync_private_state().simulate();
-
-    const notes = await pxe.getNotes({ txHash: txReceipt.txHash });
+    const notes = await pxe.getNotes({ txHash: txReceipt.txHash, contractAddress: deployedContract.address });
 
     expect(notes.length).toBe(1);
   });

yarn-project/pxe/src/contract_function_simulator/pxe_oracle_interface.test.ts

@@ -487,7 +487,7 @@ describe('PXEOracleInterface', () => {
       );
 
       // Verify note was stored
-      const notes = await noteDataProvider.getNotes({ recipient: recipient.address });
+      const notes = await noteDataProvider.getNotes({ recipient: recipient.address, contractAddress });
       expect(notes).toHaveLength(1);
       expect(notes[0].noteHash.equals(noteHash)).toBe(true);
     });
@@ -537,7 +537,7 @@ describe('PXEOracleInterface', () => {
       );
 
       // Verify note was removed
-      const notes = await noteDataProvider.getNotes({ recipient: recipient.address });
+      const notes = await noteDataProvider.getNotes({ recipient: recipient.address, contractAddress });
       expect(notes).toHaveLength(0);
     });
 
@@ -607,7 +607,11 @@ describe('PXEOracleInterface', () => {
       );
 
       // Verify note was stored and not removed
-      const notes = await noteDataProvider.getNotes({ recipient: recipient.address, status: NoteStatus.ACTIVE });
+      const notes = await noteDataProvider.getNotes({
+        recipient: recipient.address,
+        contractAddress,
+        status: NoteStatus.ACTIVE,
+      });
       expect(notes).toHaveLength(1);
       expect(notes[0].noteHash.equals(noteHash)).toBe(true);
     });
@@ -823,7 +827,7 @@ describe('PXEOracleInterface', () => {
 
     beforeEach(async () => {
       // Check that there are no notes in the database
-      const notes = await noteDataProvider.getNotes({});
+      const notes = await noteDataProvider.getNotes({ contractAddress });
       expect(notes).toHaveLength(0);
 
       // Check that the expected number of accounts is present

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

@@ -658,6 +658,9 @@ export class PXEService implements PXE {
   }
 
   public async getNotes(filter: NotesFilter): Promise<UniqueNote[]> {
+    // We need to manually trigger private state sync to have a guarantee that all the events are available.
+    await this.simulateUtility('sync_private_state', [], filter.contractAddress);
+
     const noteDaos = await this.noteDataProvider.getNotes(filter);
 
     const extendedNotes = noteDaos.map(async dao => {

yarn-project/pxe/src/storage/note_data_provider/note_data_provider.test.ts

@@ -18,32 +18,57 @@ describe('NoteDataProvider', () => {
   let storageSlots: Fr[];
   let notes: NoteDao[];
 
+  type NotesFilterWithoutContract = Omit<NotesFilter, 'contractAddress'>;
+  async function getNotesForAllContracts(filter: NotesFilterWithoutContract) {
+    const notesArrays = await Promise.all(
+      contractAddresses.map(contractAddress => noteDataProvider.getNotes({ ...filter, contractAddress })),
+    );
+    return notesArrays.flat();
+  }
+
   beforeEach(async () => {
     const store = await openTmpStore('note_data_provider_test');
     noteDataProvider = await NoteDataProvider.create(store);
   });
 
   const filteringTests: [() => Promise<NotesFilter>, () => Promise<NoteDao[]>][] = [
-    [() => Promise.resolve({}), () => Promise.resolve(notes)],
-
     [
       () => Promise.resolve({ contractAddress: contractAddresses[0] }),
       () => Promise.resolve(notes.filter(note => note.contractAddress.equals(contractAddresses[0]))),
     ],
     [async () => ({ contractAddress: await AztecAddress.random() }), () => Promise.resolve([])],
 
     [
-      () => Promise.resolve({ storageSlot: storageSlots[0] }),
-      () => Promise.resolve(notes.filter(note => note.storageSlot.equals(storageSlots[0]))),
+      () => Promise.resolve({ contractAddress: contractAddresses[0], storageSlot: storageSlots[0] }),
+      () =>
+        Promise.resolve(
+          notes.filter(
+            note => note.storageSlot.equals(storageSlots[0]) && note.contractAddress.equals(contractAddresses[0]),
+          ),
+        ),
+    ],
+    [
+      () => Promise.resolve({ contractAddress: contractAddresses[0], storageSlot: Fr.random() }),
+      () => Promise.resolve([]),
     ],
-    [() => Promise.resolve({ storageSlot: Fr.random() }), () => Promise.resolve([])],
 
-    [() => Promise.resolve({ txHash: notes[0].txHash }), () => Promise.resolve([notes[0]])],
-    [() => Promise.resolve({ txHash: randomTxHash() }), () => Promise.resolve([])],
+    [
+      () => Promise.resolve({ contractAddress: contractAddresses[0], txHash: notes[0].txHash }),
+      () => Promise.resolve([notes[0]]),
+    ],
+    [
+      () => Promise.resolve({ contractAddress: contractAddresses[0], txHash: randomTxHash() }),
+      () => Promise.resolve([]),
+    ],
 
     [
-      () => Promise.resolve({ recipient: recipients[0] }),
-      () => Promise.resolve(notes.filter(note => note.recipient.equals(recipients[0]))),
+      () => Promise.resolve({ contractAddress: contractAddresses[0], recipient: recipients[0] }),
+      () =>
+        Promise.resolve(
+          notes.filter(
+            note => note.recipient.equals(recipients[0]) && note.contractAddress.equals(contractAddresses[0]),
+          ),
+        ),
     ],
 
     [
@@ -119,8 +144,8 @@ describe('NoteDataProvider', () => {
       notesToNullify,
     );
 
-    const actualNotesWithDefault = await noteDataProvider.getNotes({});
-    const actualNotesWithActive = await noteDataProvider.getNotes({ status: NoteStatus.ACTIVE });
+    const actualNotesWithDefault = await getNotesForAllContracts({});
+    const actualNotesWithActive = await getNotesForAllContracts({ status: NoteStatus.ACTIVE });
 
     expect(actualNotesWithDefault).toEqual(actualNotesWithActive);
     expect(actualNotesWithActive).toEqual(notes.filter(note => !notesToNullify.includes(note)));
@@ -140,7 +165,7 @@ describe('NoteDataProvider', () => {
     );
     await expect(noteDataProvider.unnullifyNotesAfter(98)).resolves.toEqual(undefined);
 
-    const result = await noteDataProvider.getNotes({ status: NoteStatus.ACTIVE, recipient: recipients[0] });
+    const result = await getNotesForAllContracts({ status: NoteStatus.ACTIVE, recipient: recipients[0] });
 
     expect(result.sort()).toEqual([...notesToNullify].sort());
   });
@@ -158,7 +183,7 @@ describe('NoteDataProvider', () => {
       notesToNullify,
     );
 
-    const result = await noteDataProvider.getNotes({
+    const result = await getNotesForAllContracts({
       status: NoteStatus.ACTIVE_OR_NULLIFIED,
     });
 
@@ -172,19 +197,19 @@ describe('NoteDataProvider', () => {
 
     await noteDataProvider.addNotes(notes.slice(5), recipients[1]);
 
-    const recipient0Notes = await noteDataProvider.getNotes({
+    const recipient0Notes = await getNotesForAllContracts({
       scopes: [recipients[0]],
     });
 
     expect(recipient0Notes.sort()).toEqual(notes.slice(0, 5).sort());
 
-    const recipient1Notes = await noteDataProvider.getNotes({
+    const recipient1Notes = await getNotesForAllContracts({
       scopes: [recipients[1]],
     });
 
     expect(recipient1Notes.sort()).toEqual(notes.slice(5).sort());
 
-    const bothRecipientNotes = await noteDataProvider.getNotes({
+    const bothRecipientNotes = await getNotesForAllContracts({
       scopes: [recipients[0], recipients[1]],
     });
 
@@ -196,12 +221,12 @@ describe('NoteDataProvider', () => {
     await noteDataProvider.addNotes([notes[0]], recipients[1]);
 
     await expect(
-      noteDataProvider.getNotes({
+      getNotesForAllContracts({
         scopes: [recipients[0]],
       }),
     ).resolves.toEqual([notes[0]]);
     await expect(
-      noteDataProvider.getNotes({
+      getNotesForAllContracts({
         scopes: [recipients[1]],
       }),
     ).resolves.toEqual([notes[0]]);
@@ -219,12 +244,12 @@ describe('NoteDataProvider', () => {
     ).resolves.toEqual([notes[0]]);
 
     await expect(
-      noteDataProvider.getNotes({
+      getNotesForAllContracts({
         scopes: [recipients[0]],
       }),
     ).resolves.toEqual([]);
     await expect(
-      noteDataProvider.getNotes({
+      getNotesForAllContracts({
         scopes: [recipients[1]],
       }),
     ).resolves.toEqual([]);
@@ -234,7 +259,7 @@ describe('NoteDataProvider', () => {
     await noteDataProvider.addNotes(notes, recipients[0]);
 
     await noteDataProvider.removeNotesAfter(5);
-    const result = await noteDataProvider.getNotes({ scopes: [recipients[0]] });
+    const result = await getNotesForAllContracts({ scopes: [recipients[0]] });
     expect(new Set(result)).toEqual(new Set(notes.slice(0, 6)));
   });
 });

yarn-project/pxe/src/storage/note_data_provider/note_data_provider.ts

@@ -202,19 +202,17 @@ export class NoteDataProvider implements DataProvider {
             ? await toArray(
                 this.#notesByTxHashAndScope.get(formattedScopeString)!.getValuesAsync(filter.txHash.toString()),
               )
-            : filter.contractAddress
+            : filter.storageSlot
               ? await toArray(
+                  this.#notesByStorageSlotAndScope
+                    .get(formattedScopeString)!
+                    .getValuesAsync(filter.storageSlot.toString()),
+                )
+              : await toArray(
                   this.#notesByContractAndScope
                     .get(formattedScopeString)!
                     .getValuesAsync(filter.contractAddress.toString()),
-                )
-              : filter.storageSlot
-                ? await toArray(
-                    this.#notesByStorageSlotAndScope
-                      .get(formattedScopeString)!
-                      .getValuesAsync(filter.storageSlot.toString()),
-                  )
-                : await toArray(this.#notesByRecipientAndScope.get(formattedScopeString)!.valuesAsync()),
+                ),
       );
     }
 
@@ -229,11 +227,9 @@ export class NoteDataProvider implements DataProvider {
           ? await toArray(this.#nullifiedNotesByRecipient.getValuesAsync(filter.recipient.toString()))
           : filter.txHash
             ? await toArray(this.#nullifiedNotesByTxHash.getValuesAsync(filter.txHash.toString()))
-            : filter.contractAddress
-              ? await toArray(this.#nullifiedNotesByContract.getValuesAsync(filter.contractAddress.toString()))
-              : filter.storageSlot
-                ? await toArray(this.#nullifiedNotesByStorageSlot.getValuesAsync(filter.storageSlot.toString()))
-                : await toArray(this.#nullifiedNotes.keysAsync()),
+            : filter.storageSlot
+              ? await toArray(this.#nullifiedNotesByStorageSlot.getValuesAsync(filter.storageSlot.toString()))
+              : await toArray(this.#nullifiedNotesByContract.getValuesAsync(filter.contractAddress.toString())),
         notes: this.#nullifiedNotes,
       });
     }
@@ -247,7 +243,7 @@ export class NoteDataProvider implements DataProvider {
         }
 
         const note = NoteDao.fromBuffer(serializedNote);
-        if (filter.contractAddress && !note.contractAddress.equals(filter.contractAddress)) {
+        if (!note.contractAddress.equals(filter.contractAddress)) {
           continue;
         }
 
@@ -334,6 +330,21 @@ export class NoteDataProvider implements DataProvider {
   }
 
   async getSize() {
-    return (await this.getNotes({})).reduce((sum, note) => sum + note.getSize(), 0);
+    const scopes = await toArray(this.#scopes.keysAsync());
+    const contractAddresses = new Set<string>();
+
+    // Collect all unique contract addresses across all scopes
+    for (const scope of scopes) {
+      const addresses = await toArray(this.#notesByContractAndScope.get(scope)!.keysAsync());
+      addresses.forEach(addr => contractAddresses.add(addr));
+    }
+
+    // Get all notes for each contract address
+    const allNotes = await Promise.all(
+      Array.from(contractAddresses).map(addr => this.getNotes({ contractAddress: AztecAddress.fromString(addr) })),
+    );
+
+    // Reduce all notes to get total size
+    return allNotes.flat().reduce((sum, note) => sum + note.getSize(), 0);
   }
 }