feat: store partial notes in pxe

noir-projects/aztec-nr/aztec/src/note/lifecycle.nr

@@ -1,6 +1,6 @@
 use crate::context::{PrivateContext, PublicContext};
 use crate::note::{
-    note_header::NoteHeader, note_interface::NoteInterface,
+    note_header::NoteHeader, note_interface::{NoteInterface, CompletableNoteInterface},
     utils::{compute_note_hash_for_insertion, compute_note_hash_for_consumption, compute_partial_note_hash}
 };
 use crate::oracle::notes::{notify_created_note, notify_nullified_note};
@@ -43,7 +43,7 @@ pub fn create_partial_note<Note, N>(
     storage_slot: Field,
     note: &mut Note,
     broadcast: bool
-) -> Field where Note: NoteInterface<N> {
+) -> Field where Note: NoteInterface<N> + CompletableNoteInterface {
     let contract_address = (*context).this_address();
 
     let mut header = NoteHeader { contract_address, storage_slot, nonce: 0, is_transient: true, partial_note_hash: 0 };
@@ -57,7 +57,7 @@ pub fn create_partial_note<Note, N>(
     Note::set_header(note, header);
 
     if broadcast {
-        Note::broadcast(*note, context, storage_slot);
+        CompletableNoteInterface::broadcast_partial_data(*note, context, storage_slot);
     }
 
     partial_note_hash
@@ -66,8 +66,9 @@ pub fn create_partial_note<Note, N>(
 pub fn complete_partial_note_from_public<Note, N>(
     context: &mut PublicContext,
     partial_note_hash: Field,
-    note: &mut Note
-) where Note: NoteInterface<N> {
+    note: &mut Note,
+    broadcast: bool
+) where Note: NoteInterface<N> + CompletableNoteInterface {
     let contract_address = (*context).this_address();
 
     let header = NoteHeader { contract_address, storage_slot: 0, nonce: 0, is_transient: true, partial_note_hash };
@@ -76,10 +77,11 @@ pub fn complete_partial_note_from_public<Note, N>(
     // As `is_transient` is true, this will compute the inner note hsah
     let note_hash = compute_note_hash_for_insertion(*note);
 
-    // No need to broadcast.
-    // Whoever created the partial note "should" already have the missing info ncessary to complete the note in their PXE.
-    // Also broadcsating would require the note's storage slot.
     context.push_new_note_hash(note_hash);
+
+    if broadcast {
+        CompletableNoteInterface::broadcast_completion_from_public(*note, context);
+    }
 }
 
 pub fn create_note_hash_from_public<Note, N>(

noir-projects/aztec-nr/aztec/src/note/note_header.nr

@@ -23,8 +23,4 @@ impl NoteHeader {
     pub fn new(contract_address: AztecAddress, nonce: Field, storage_slot: Field) -> Self {
         NoteHeader { contract_address, nonce, storage_slot, is_transient: false, partial_note_hash: 0 }
     }
-
-    pub fn is_partial(self) -> bool {
-        self.partial_note_hash != 0
-    }
 }

noir-projects/aztec-nr/aztec/src/note/note_interface.nr

@@ -1,4 +1,4 @@
-use crate::context::PrivateContext;
+use crate::context::{PrivateContext, PublicContext};
 use crate::note::note_header::NoteHeader;
 
 // docs:start:note_interface
@@ -24,3 +24,8 @@ trait NoteInterface<N> {
     fn get_note_type_id() -> Field;
 }
 // docs:end:note_interface
+
+trait CompletableNoteInterface {
+    fn broadcast_partial_data(self, context: &mut PrivateContext, slot: Field) -> ();
+    fn broadcast_completion_from_public(self, context: &mut PublicContext) -> ();
+}

noir-projects/aztec-nr/aztec/src/state_vars/private_set.nr

@@ -7,8 +7,8 @@ use crate::context::{PrivateContext, PublicContext, Context};
 use crate::note::{
     lifecycle::{create_note, create_partial_note, create_note_hash_from_public, destroy_note},
     note_getter::{get_notes, view_notes}, note_getter_options::NoteGetterOptions,
-    note_header::NoteHeader, note_interface::NoteInterface, note_viewer_options::NoteViewerOptions,
-    utils::compute_note_hash_for_consumption
+    note_header::NoteHeader, note_interface::{NoteInterface, CompletableNoteInterface},
+    note_viewer_options::NoteViewerOptions, utils::compute_note_hash_for_consumption
 };
 use crate::state_vars::storage::Storage;
 
@@ -43,7 +43,7 @@ impl<Note> PrivateSet<Note> {
         self,
         note: &mut Note,
         broadcast: bool
-    ) -> Field where Note: NoteInterface<N> {
+    ) -> Field where Note: NoteInterface<N> + CompletableNoteInterface {
         create_partial_note(
             self.context.private.unwrap(),
             self.storage_slot,

noir-projects/noir-contracts/contracts/token_contract/src/main.nr

@@ -399,11 +399,8 @@ contract Token {
 
         storage.balances.sub(from, U128::from_integer(amount));
 
-        let partial_note_0 = storage.balances.add_partial(from);
-        let partial_note_1 = storage.balances.add_partial(to);
-
-        let partial_note_hash_0 = partial_note_0.get_header().partial_note_hash;
-        let partial_note_hash_1 = partial_note_1.get_header().partial_note_hash;
+        let partial_note_hash_0 = storage.balances.add_partial(from);
+        let partial_note_hash_1 = storage.balances.add_partial(to);
 
         let hash = compute_partial_notes_pair_hash(
             partial_note_hash_0,
@@ -418,18 +415,6 @@ contract Token {
             [hash.to_field()]
         );
 
-        // HACK: emit private data as unencrypted logs until PXE can track partial notes
-        // (the partial note is emitted as an encrypted log but the PXE drops it because its not in the tree)
-        emit_unencrypted_log_from_private(
-            &mut context,
-            [partial_note_0.randomness, partial_note_1.randomness]
-        );
-
-        emit_unencrypted_log_from_private(
-            &mut context,
-            [partial_note_0.get_header().storage_slot, partial_note_1.get_header().storage_slot]
-        );
-
         // emit as unencrypted for e2e testing as return values are not accessible there
         emit_unencrypted_log_from_private(&mut context, [partial_note_hash_0, partial_note_hash_1]);
 
@@ -455,12 +440,10 @@ contract Token {
         context.push_new_nullifier(hash.to_field(), 0);
 
         let mut token_note_0 = TokenNote::new_from_public(U128::from_integer(amounts[0]));
-        complete_partial_note_from_public(&mut context, partial_note_hashes[0], &mut token_note_0);
+        complete_partial_note_from_public(&mut context, partial_note_hashes[0], &mut token_note_0, true);
 
         let mut token_note_1 = TokenNote::new_from_public(U128::from_integer(amounts[1]));
-        complete_partial_note_from_public(&mut context, partial_note_hashes[1], &mut token_note_1);
-
-        emit_unencrypted_log(&mut context, amounts);
+        complete_partial_note_from_public(&mut context, partial_note_hashes[1], &mut token_note_1, true);
     }
 }
 // docs:end:token_all

noir-projects/noir-contracts/contracts/token_contract/src/types/balances_map.nr

@@ -5,7 +5,8 @@ use dep::aztec::{
     state_vars::{PrivateSet, Map},
     note::{
     note_getter::view_notes, note_getter_options::{NoteGetterOptions, SortOrder},
-    note_viewer_options::NoteViewerOptions, note_header::NoteHeader, note_interface::NoteInterface
+    note_viewer_options::NoteViewerOptions, note_header::NoteHeader,
+    note_interface::{NoteInterface, CompletableNoteInterface}
 }
 };
 use crate::types::token_note::{TokenNote, OwnedNote};
@@ -53,11 +54,13 @@ impl<T> BalancesMap<T> {
         balance
     }
 
-    pub fn add_partial<T_SERIALIZED_LEN>(self: Self, owner: AztecAddress) -> T where T: NoteInterface<T_SERIALIZED_LEN> + OwnedNote {
+    pub fn add_partial<T_SERIALIZED_LEN>(
+        self: Self,
+        owner: AztecAddress
+    ) -> Field where T: NoteInterface<T_SERIALIZED_LEN> + CompletableNoteInterface + OwnedNote {
         let mut addend_note = T::new(U128::empty(), owner);
-        let _partial_note_hash = self.map.at(owner).insert_partial(&mut addend_note, true);
-        // TODO once PXE tracks partial notes, just return the hash instead of the whole note
-        addend_note
+        let partial_note_hash = self.map.at(owner).insert_partial(&mut addend_note, true);
+        partial_note_hash
     }
 
     pub fn add<T_SERIALIZED_LEN>(

noir-projects/noir-contracts/contracts/token_contract/src/types/token_note.nr

@@ -1,7 +1,11 @@
 use dep::aztec::{
-    protocol_types::address::AztecAddress,
-    note::{note_header::NoteHeader, note_interface::NoteInterface, utils::compute_note_hash_for_consumption},
-    context::PrivateContext, log::emit_encrypted_log, hash::pedersen_hash
+    protocol_types::{address::AztecAddress, traits::is_empty},
+    note::{
+    note_header::NoteHeader, note_interface::{NoteInterface, CompletableNoteInterface},
+    utils::compute_note_hash_for_consumption
+},
+    context::{PrivateContext, PublicContext}, log::{emit_encrypted_log, emit_unencrypted_log},
+    hash::pedersen_hash
 };
 use dep::aztec::oracle::{rand::rand, nullifier_key::get_nullifier_secret_key, get_public_key::get_public_key};
 use dep::std::option::Option;
@@ -88,7 +92,7 @@ impl NoteInterface<TOKEN_NOTE_LEN> for TokenNote {
     // Broadcasts the note as an encrypted log on L1.
     fn broadcast(self, context: &mut PrivateContext, slot: Field) {
       // We only bother inserting the note if non-empty to save funds on gas.
-      if !(self.amount == U128::from_integer(0)) | self.header.is_partial() {
+      if !is_empty(self.amount) {
           let encryption_pub_key = get_public_key(self.owner);
           emit_encrypted_log(
               context,
@@ -127,6 +131,31 @@ impl OwnedNote for TokenNote {
     }
 }
 
+impl CompletableNoteInterface for TokenNote {
+    fn broadcast_partial_data(self, context: &mut PrivateContext, slot: Field) {
+        let encryption_pub_key = get_public_key(self.owner);
+
+        emit_encrypted_log(
+            context,
+            (*context).this_address(),
+            slot,
+            TokenNote::get_note_type_id(),
+            encryption_pub_key,
+            self.serialize_content(),
+        );
+    }
+
+    fn broadcast_completion_from_public(self, context: &mut PublicContext) {
+        if !is_empty(self.amount) {
+            // TODO change this with a well-structured Event so the PXE can listen for them
+            emit_unencrypted_log(
+                context,
+                [TokenNote::get_note_type_id(), self.amount.to_field()]
+            );
+        }
+    }
+}
+
 impl TokenNote {
     pub fn new_from_public(amount: U128) -> Self {
         Self { amount, owner: AztecAddress::empty(), randomness: 0, header: NoteHeader::empty() }

yarn-project/aztec.js/src/wallet/base_wallet.ts

@@ -126,4 +126,13 @@ export abstract class BaseWallet implements Wallet {
   isContractClassPubliclyRegistered(id: Fr): Promise<boolean> {
     return this.pxe.isContractClassPubliclyRegistered(id);
   }
+
+  completePartialNote(
+    contractAddress: AztecAddress,
+    noteTypeId: Fr,
+    patches: [number | Fr, Fr][],
+    txHash: TxHash,
+  ): Promise<void> {
+    return this.pxe.completePartialNote(contractAddress, noteTypeId, patches, txHash);
+  }
 }

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

@@ -287,5 +287,20 @@ export interface PXE {
    * @param id - Identifier of the class.
    */
   isContractClassPubliclyRegistered(id: Fr): Promise<boolean>;
+
+  /**
+   * Attempts to complete a partial note.
+   *
+   * @param contractAddress - The contract whose note is to be completed.
+   * @param noteTypeId - The type of the note to be completed.
+   * @param patches - The patches to apply to the serialized partial note data
+   * @param completionTxHash - The hash of the transaction that completed the partial note.
+   */
+  completePartialNote(
+    contractAddress: AztecAddress,
+    noteTypeId: Fr,
+    patches: [number | Fr, Fr][],
+    completionTxHash: TxHash,
+  ): Promise<void>;
 }
 // docs:end:pxe-interface

yarn-project/circuit-types/src/stats/stats.ts

@@ -111,8 +111,12 @@ export type NoteProcessorStats = {
   seen: number;
   /** How many notes had decryption deferred due to a missing contract */
   deferred: number;
+  /** How many partial notes we've discovered so far */
+  partial: number;
   /** How many notes were successfully decrypted. */
   decrypted: number;
+  /** How many partial notes were completed */
+  completed: number;
   /** How many notes failed processing. */
   failed: number;
   /** How many blocks were spanned.  */