remove deferred notes from pxe db

compute nullifiers for deferred notes when decoded

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

@@ -292,7 +292,7 @@ describe('e2e_2_pxes', () => {
 
     await expectsNumOfEncryptedLogsInTheLastBlockToBe(aztecNode, 1);
 
-    // // Transfer funds from A to B via PXE A
+    // Transfer funds from A to B via PXE A
     const contractWithWalletA = await TokenContract.at(tokenAddress, walletA);
     const receiptAToB = await contractWithWalletA.methods
       .transfer(userA.address, userB.address, transferAmount1, 0)
@@ -311,4 +311,73 @@ describe('e2e_2_pxes', () => {
     await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance - transferAmount1);
     await expectTokenBalance(walletB, tokenAddress, userB.address, transferAmount1);
   });
+
+  it('permits sending funds to a user, and spending them, before they have registered the contract', async () => {
+    const initialBalance = 987n;
+    const transferAmount1 = 654n;
+    const transferAmount2 = 323n;
+
+    // setup an account that is shared across PXEs
+    const sharedPrivateKey = GrumpkinScalar.random();
+    const sharedAccountOnA = getUnsafeSchnorrAccount(pxeA, sharedPrivateKey, Fr.random());
+    const sharedAccountAddress = sharedAccountOnA.getCompleteAddress();
+    const sharedWalletOnA = await sharedAccountOnA.waitDeploy();
+    await expect(sharedWalletOnA.isAccountStateSynchronized(sharedAccountAddress.address)).resolves.toBe(true);
+
+    const sharedAccountOnB = getUnsafeSchnorrAccount(pxeB, sharedPrivateKey, sharedAccountAddress);
+    await sharedAccountOnB.register();
+    const sharedWalletOnB = await sharedAccountOnB.getWallet();
+
+    await pxeA.registerRecipient(userB);
+
+    // deploy the contract on PXE A
+    const completeTokenAddress = await deployTokenContract(initialBalance, userA.address, pxeA);
+    const tokenAddress = completeTokenAddress.address;
+
+    // Transfer funds from A to Shared Wallet via PXE A
+    const contractWithWalletA = await TokenContract.at(tokenAddress, walletA);
+    const receiptAToShared = await contractWithWalletA.methods
+      .transfer(userA.address, sharedAccountAddress.address, transferAmount1, 0)
+      .send()
+      .wait();
+    expect(receiptAToShared.status).toBe(TxStatus.MINED);
+
+    // Now send funds from Shared Wallet to B via PXE A
+    const contractWithSharedWalletA = await TokenContract.at(tokenAddress, sharedWalletOnA);
+    const receiptSharedToB = await contractWithSharedWalletA.methods
+      .transfer(sharedAccountAddress.address, userB.address, transferAmount2, 0)
+      .send()
+      .wait();
+    expect(receiptSharedToB.status).toBe(TxStatus.MINED);
+
+    // check balances from PXE-A's perspective
+    await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance - transferAmount1);
+    await expectTokenBalance(
+      sharedWalletOnA,
+      tokenAddress,
+      sharedAccountAddress.address,
+      transferAmount1 - transferAmount2,
+    );
+
+    // now add the contract and check balances from PXE-B's perspective.
+    // The process should be:
+    // PXE-B had previously deferred the notes from A -> Shared, and Shared -> B
+    // PXE-B adds the contract
+    // PXE-B reprocesses the deferred notes, and sees the nullifier for A -> Shared
+    await pxeB.addContracts([
+      {
+        artifact: TokenContract.artifact,
+        completeAddress: completeTokenAddress,
+        portalContract: EthAddress.ZERO,
+      },
+    ]);
+    await expectTokenBalance(walletB, tokenAddress, userB.address, transferAmount2);
+    await expect(sharedWalletOnB.isAccountStateSynchronized(sharedAccountAddress.address)).resolves.toBe(true);
+    await expectTokenBalance(
+      sharedWalletOnB,
+      tokenAddress,
+      sharedAccountAddress.address,
+      transferAmount1 - transferAmount2,
+    );
+  });
 });

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

@@ -101,10 +101,10 @@ export class KVPxeDatabase implements PxeDatabase {
     }
   }
 
-  async addDeferredNotes(notes: DeferredNoteDao[]): Promise<void> {
-    const newLength = await this.#deferredNotes.push(...notes.map(note => note.toBuffer()));
-    for (const [index, note] of notes.entries()) {
-      const noteId = newLength - notes.length + index;
+  async addDeferredNotes(deferredNotes: DeferredNoteDao[]): Promise<void> {
+    const newLength = await this.#deferredNotes.push(...deferredNotes.map(note => note.toBuffer()));
+    for (const [index, note] of deferredNotes.entries()) {
+      const noteId = newLength - deferredNotes.length + index;
       await this.#deferredNotesByContract.set(note.contractAddress.toString(), noteId);
     }
   }
@@ -125,6 +125,36 @@ export class KVPxeDatabase implements PxeDatabase {
     return Promise.resolve(notes);
   }
 
+  /**
+   * Removes all deferred notes for a given contract address.
+   * @param contractAddress - the contract address to remove deferred notes for
+   * @returns an array of the removed deferred notes
+   *
+   * @remarks We only remove indices from the deferred notes by contract map, but not the actual deferred notes.
+   * This is safe because our only getter for deferred notes is by contract address.
+   * If we should add a more general getter, we will need a delete vector for deferred notes as well,
+   * analogous to this.#nullifiedNotes.
+   */
+  removeDeferredNotesByContract(contractAddress: AztecAddress): Promise<DeferredNoteDao[]> {
+    return this.#db.transaction(() => {
+      const deferredNotes: DeferredNoteDao[] = [];
+      const indices = this.#deferredNotesByContract.getValues(contractAddress.toString());
+
+      for (const index of indices) {
+        const deferredNoteBuffer = this.#deferredNotes.at(index);
+        if (!deferredNoteBuffer) {
+          continue;
+        } else {
+          deferredNotes.push(DeferredNoteDao.fromBuffer(deferredNoteBuffer));
+        }
+
+        void this.#deferredNotesByContract.deleteValue(contractAddress.toString(), index);
+      }
+
+      return deferredNotes;
+    });
+  }
+
   *#getAllNonNullifiedNotes(): IterableIterator<NoteDao> {
     for (const [index, serialized] of this.#notes.entries()) {
       if (this.#nullifiedNotes.has(index)) {
@@ -185,6 +215,9 @@ export class KVPxeDatabase implements PxeDatabase {
   }
 
   removeNullifiedNotes(nullifiers: Fr[], account: PublicKey): Promise<NoteDao[]> {
+    if (nullifiers.length === 0) {
+      return Promise.resolve([]);
+    }
     const nullifierSet = new Set(nullifiers.map(n => n.toString()));
     return this.#db.transaction(() => {
       const notesIds = this.#notesByOwner.getValues(account.toString());

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

@@ -17,6 +17,7 @@ import { PxeDatabase } from './pxe_database.js';
  */
 export class MemoryDB extends MemoryContractDatabase implements PxeDatabase {
   private notesTable: NoteDao[] = [];
+
   private treeRoots: Record<MerkleTreeId, Fr> | undefined;
   private globalVariablesHash: Fr | undefined;
   private blockNumber: number | undefined;
@@ -65,6 +66,11 @@ export class MemoryDB extends MemoryContractDatabase implements PxeDatabase {
     throw new Error('Method not implemented.');
   }
 
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  public removeDeferredNotesByContract(contractAddress: AztecAddress): Promise<DeferredNoteDao[]> {
+    throw new Error('Method not implemented.');
+  }
+
   public addCapsule(capsule: Fr[]): Promise<void> {
     this.capsuleStack.push(capsule);
     return Promise.resolve();

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

@@ -73,6 +73,12 @@ export interface PxeDatabase extends ContractDatabase {
    */
   getDeferredNotesByContract(contractAddress: AztecAddress): Promise<DeferredNoteDao[]>;
 
+  /**
+   * Remove deferred notes for a given contract address.
+   * @param contractAddress - The contract address to remove the deferred notes for.
+   */
+  removeDeferredNotesByContract(contractAddress: AztecAddress): Promise<DeferredNoteDao[]>;
+
   /**
    * Remove nullified notes associated with the given account and nullifiers.
    *

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

@@ -166,7 +166,7 @@ export class NoteProcessor {
                   deferredNoteDaos.push(deferredNoteDao);
                 } else {
                   this.stats.failed++;
-                  this.log.warn(`Could not process note because of "${e}". Skipping note...`);
+                  this.log.warn(`Could not process note because of "${e}". Discarding note...`);
                 }
               }
             }
@@ -241,11 +241,15 @@ export class NoteProcessor {
   }
 
   /**
-   * Retry processing the given deferred notes because we now have the contract code.
+   * Retry decoding the given deferred notes because we now have the contract code.
    *
    * @param deferredNoteDaos - notes that we have previously deferred because the contract was not found
+   * @returns An array of NoteDaos that were successfully decoded.
+   *
+   * @remarks Caller is responsible for making sure that we have the contract for the
+   * deferred notes provided: we will not retry notes that fail again.
    */
-  public async retryDeferredNotes(deferredNoteDaos: DeferredNoteDao[]) {
+  public async decodeDeferredNotes(deferredNoteDaos: DeferredNoteDao[]): Promise<NoteDao[]> {
     const excludedIndices: Set<number> = new Set();
     const noteDaos: NoteDao[] = [];
     for (const deferredNote of deferredNoteDaos) {
@@ -268,22 +272,10 @@ export class NoteProcessor {
         this.stats.decrypted++;
       } catch (e) {
         this.stats.failed++;
-        this.log.warn(`Could not process deferred note because of "${e}". Skipping note...`);
+        this.log.warn(`Could not process deferred note because of "${e}". Discarding note...`);
       }
     }
 
-    if (noteDaos.length) {
-      await this.db.addNotes(noteDaos);
-      noteDaos.forEach(noteDao => {
-        this.log(
-          `Decoded and added deferred note for contract ${noteDao.contractAddress} at slot ${
-            noteDao.storageSlot
-          } with nullifier ${noteDao.siloedNullifier.toString()}`,
-        );
-      });
-
-      // TODO: Remove deferred notes from the database.
-      // TODO: keep track of the oldest deferred note that has been decoded, then reprocess nullifiers from that block onwards.
-    }
+    return noteDaos;
   }
 }

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

@@ -57,7 +57,6 @@ import { TxPXEProcessingStats } from '@aztec/types/stats';
 
 import { PXEServiceConfig, getPackageInfo } from '../config/index.js';
 import { ContractDataOracle } from '../contract_data_oracle/index.js';
-import { DeferredNoteDao } from '../database/deferred_note_dao.js';
 import { PxeDatabase } from '../database/index.js';
 import { NoteDao } from '../database/note_dao.js';
 import { KernelOracle } from '../kernel_oracle/index.js';
@@ -210,7 +209,7 @@ export class PXEService implements PXE {
       const portalInfo =
         contract.portalContract && !contract.portalContract.isZero() ? ` with portal ${contract.portalContract}` : '';
       this.log.info(`Added contract ${contract.name} at ${contractAztecAddress}${portalInfo}`);
-      await this.synchronizer.retryDeferredNotesForContract(contractAztecAddress);
+      await this.synchronizer.reprocessDeferredNotesForContract(contractAztecAddress);
     }
   }
 
@@ -491,10 +490,6 @@ export class PXEService implements PXE {
     return nodeInfo;
   }
 
-  #retryDeferredNote(deferredNote: DeferredNoteDao) {
-    this.synchronizer;
-  }
-
   /**
    * Retrieves the simulation parameters required to run an ACIR simulation.
    * This includes the contract address, function artifact, portal contract address, and historical tree roots.

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

@@ -10,12 +10,14 @@ import {
   L2BlockContext,
   L2BlockL2Logs,
   LogType,
+  MerkleTreeId,
   TxHash,
 } from '@aztec/types';
 import { NoteProcessorCaughtUpStats } from '@aztec/types/stats';
 
 import { DeferredNoteDao } from '../database/deferred_note_dao.js';
 import { PxeDatabase } from '../database/index.js';
+import { NoteDao } from '../database/note_dao.js';
 import { NoteProcessor } from '../note_processor/index.js';
 
 /**
@@ -338,7 +340,7 @@ export class Synchronizer {
    * Retry decoding any deferred notes for the specified contract address.
    * @param contractAddress - the contract address that has just been added
    */
-  public async retryDeferredNotesForContract(contractAddress: AztecAddress) {
+  public async reprocessDeferredNotesForContract(contractAddress: AztecAddress) {
     const deferredNotes = await this.db.getDeferredNotesByContract(contractAddress);
 
     // group deferred notes by txHash to properly deal with possible duplicates
@@ -349,12 +351,51 @@ export class Synchronizer {
       txHashToDeferredNotes.set(note.txHash, notesForTx);
     }
 
+    // keep track of decoded notes
+    const newNotes: NoteDao[] = [];
     // now process each txHash
     for (const deferredNotes of txHashToDeferredNotes.values()) {
       // to be safe, try each note processor in case the deferred notes are for different accounts.
       for (const processor of this.noteProcessors) {
-        await processor.retryDeferredNotes(deferredNotes.filter(n => n.publicKey.equals(processor.publicKey)));
+        const decodedNotes = await processor.decodeDeferredNotes(
+          deferredNotes.filter(n => n.publicKey.equals(processor.publicKey)),
+        );
+        newNotes.push(...decodedNotes);
       }
     }
+
+    // now drop the deferred notes, and add the decoded notes
+    await this.db.removeDeferredNotesByContract(contractAddress);
+    await this.db.addNotes(newNotes);
+
+    newNotes.forEach(noteDao => {
+      this.log(
+        `Decoded deferred note for contract ${noteDao.contractAddress} at slot ${
+          noteDao.storageSlot
+        } with nullifier ${noteDao.siloedNullifier.toString()}`,
+      );
+    });
+
+    // now group the decoded notes by public key
+    const publicKeyToNotes: Map<PublicKey, NoteDao[]> = new Map();
+    for (const noteDao of newNotes) {
+      const notesForPublicKey = publicKeyToNotes.get(noteDao.publicKey) ?? [];
+      notesForPublicKey.push(noteDao);
+      publicKeyToNotes.set(noteDao.publicKey, notesForPublicKey);
+    }
+
+    // now for each group, look for the nullifiers in the nullifier tree
+    for (const [publicKey, notes] of publicKeyToNotes.entries()) {
+      const nullifiers = notes.map(n => n.siloedNullifier);
+      const relevantNullifiers: Fr[] = [];
+      for (const nullifier of nullifiers) {
+        // NOTE: this leaks information about the nullifiers I'm interested in to the node.
+        const found = await this.node.findLeafIndex('latest', MerkleTreeId.NULLIFIER_TREE, nullifier);
+        if (found) {
+          relevantNullifiers.push(nullifier);
+        }
+      }
+      await this.db.removeNullifiedNotes(relevantNullifiers, publicKey);
+    }
   }
 }