feat: partial note handling in aztec-nr (#12122)

We're back! Reopening
AztecProtocol/aztec-packages#11641, which was
reverted in AztecProtocol/aztec-packages#11797.

---------

Co-authored-by: Jan Beneš <janbenes1234@gmail.com>

Author: 2 people

aztec/src/discovery/mod.nr

@@ -0,0 +1,77 @@
+use dep::protocol_types::{
+    address::AztecAddress, constants::PRIVATE_LOG_SIZE_IN_FIELDS, debug_log::debug_log,
+};
+
+pub mod private_logs;
+pub mod partial_notes;
+pub mod nonce_discovery;
+
+/// We reserve two fields in the note private log that are not part of the note content: one for the storage slot, and
+/// one for the combined log and note type ID.
+global NOTE_PRIVATE_LOG_RESERVED_FIELDS: u32 = 2;
+
+/// The maximum length of the packed representation of a note's contents. This is limited by private log size and extra
+/// fields in the log (e.g. the combined log and note type ID).
+// TODO (#11634): we're assuming here that the entire log is plaintext, which is not true due to headers, encryption
+// padding, etc. Notes can't actually be this large.
+pub global MAX_NOTE_PACKED_LEN: u32 = PRIVATE_LOG_SIZE_IN_FIELDS - NOTE_PRIVATE_LOG_RESERVED_FIELDS;
+
+pub struct NoteHashAndNullifier {
+    /// The result of NoteInterface::compute_note_hash
+    pub note_hash: Field,
+    /// The result of NullifiableNote::compute_nullifier_without_context
+    pub inner_nullifier: Field,
+}
+
+/// A function which takes a note's packed content, address of the emitting contract, nonce, storage slot and note type
+/// ID and attempts to compute its note hash (not siloed by nonce nor address) and inner nullifier (not siloed by
+/// address).
+///
+/// This function must be user-provided as its implementation requires knowledge of how note type IDs are allocated in a
+/// contract. A typical implementation would look like this:
+///
+/// ```
+/// |packed_note_content, contract_address, nonce, storage_slot, note_type_id| {
+///     if note_type_id == MyNoteType::get_note_type_id() {
+///         assert(packed_note_content.len() == MY_NOTE_TYPE_SERIALIZATION_LENGTH);
+///         let hashes = dep::aztec::note::utils::compute_note_hash_and_optionally_a_nullifier(
+///             MyNoteType::unpack_content,
+///             note_header,
+///             true,
+///             packed_note_content.storage(),
+///         )
+///
+///         Option::some(dep::aztec::oracle::management::NoteHashesAndNullifier {
+///             note_hash: hashes[0],
+///             inner_nullifier: hashes[3],
+///         })
+///     } else if note_type_id == MyOtherNoteType::get_note_type_id() {
+///           ... // Similar to above but calling MyOtherNoteType::unpack_content
+///     } else {
+///         Option::none() // Unknown note type ID
+///     };
+/// }
+/// ```
+type ComputeNoteHashAndNullifier<Env> = fn[Env](/* packed_note_content */BoundedVec<Field, MAX_NOTE_PACKED_LEN>, /* contract_address */ AztecAddress, /* nonce */ Field, /* storage_slot */ Field, /* note_type_id */ Field) -> Option<NoteHashAndNullifier>;
+
+/// Performs the note discovery process, in which private and public logs are downloaded and inspected to find private
+/// notes, partial notes, and their completion. This is the mechanism via which PXE learns of new notes.
+///
+/// Receives the address of the contract on which discovery is performed (i.e. the contract that emitted the notes)
+/// along with its `compute_note_hash_and_nullifier` function.
+pub unconstrained fn discover_new_notes<Env>(
+    contract_address: AztecAddress,
+    compute_note_hash_and_nullifier: ComputeNoteHashAndNullifier<Env>,
+) {
+    debug_log("Performing note discovery");
+
+    private_logs::fetch_and_process_private_tagged_logs(
+        contract_address,
+        compute_note_hash_and_nullifier,
+    );
+
+    partial_notes::fetch_and_process_public_partial_note_completion_logs(
+        contract_address,
+        compute_note_hash_and_nullifier,
+    );
+}

aztec/src/discovery/nonce_discovery.nr

@@ -0,0 +1,94 @@
+use crate::{discovery::{MAX_NOTE_PACKED_LEN, NoteHashAndNullifier}, utils::array};
+
+use dep::protocol_types::{
+    address::AztecAddress,
+    constants::MAX_NOTE_HASHES_PER_TX,
+    debug_log::debug_log_format,
+    hash::{compute_note_hash_nonce, compute_siloed_note_hash, compute_unique_note_hash},
+    traits::ToField,
+};
+
+/// A struct with the discovered information of a complete note, required for delivery to PXE. Note that this is *not*
+/// the complete note information, since it does not include content, storage slot, etc.
+pub struct DiscoveredNoteInfo {
+    pub nonce: Field,
+    pub note_hash: Field,
+    pub inner_nullifier: Field,
+}
+
+/// Searches for note nonces that will result in a note that was emitted in a transaction. While rare, it is possible
+/// for multiple notes to have the exact same packed content and storage slot but different nonces, resulting in
+/// different unique note hashes. Because of this this function returns a *vector* of discovered notes, though in most
+/// cases it will contain a single element.
+///
+/// Due to how nonces are computed, this function requires knowledge of the transaction in which the note was created,
+/// more specifically the list of all unique note hashes in it plus the value of its first nullifier.
+pub unconstrained fn attempt_note_nonce_discovery<Env>(
+    unique_note_hashes_in_tx: BoundedVec<Field, MAX_NOTE_HASHES_PER_TX>,
+    first_nullifier_in_tx: Field,
+    compute_note_hash_and_nullifier: fn[Env](BoundedVec<Field, MAX_NOTE_PACKED_LEN>, AztecAddress, Field, Field, Field) -> Option<NoteHashAndNullifier>,
+    contract_address: AztecAddress,
+    storage_slot: Field,
+    note_type_id: Field,
+    packed_note_content: BoundedVec<Field, MAX_NOTE_PACKED_LEN>,
+) -> BoundedVec<DiscoveredNoteInfo, MAX_NOTE_HASHES_PER_TX> {
+    let discovered_notes = &mut BoundedVec::new();
+
+    debug_log_format(
+        "Attempting note discovery on {0} potential notes on contract {1} for storage slot {2}",
+        [unique_note_hashes_in_tx.len() as Field, contract_address.to_field(), storage_slot],
+    );
+
+    // We need to find nonces (typically just one) that result in a note hash that, once siloed into a unique note hash,
+    // is one of the note hashes created by the transaction.
+    array::for_each_in_bounded_vec(
+        unique_note_hashes_in_tx,
+        |expected_unique_note_hash, i| {
+            // Nonces are computed by hashing the first nullifier in the transaction with the index of the note in the
+            // new note hashes array. We therefore know for each note in every transaction what its nonce is.
+            let candidate_nonce = compute_note_hash_nonce(first_nullifier_in_tx, i);
+
+            // Given nonce, note content and metadata, we can compute the note hash and silo it to check if it matches
+            // the note hash at the array index we're currently processing.
+            // TODO(#11157): handle failed note_hash_and_nullifier computation
+            let hashes = compute_note_hash_and_nullifier(
+                packed_note_content,
+                contract_address,
+                candidate_nonce,
+                storage_slot,
+                note_type_id,
+            )
+                .expect(f"Failed to compute a note hash for note type {note_type_id}");
+
+            let siloed_note_hash = compute_siloed_note_hash(contract_address, hashes.note_hash);
+            let unique_note_hash = compute_unique_note_hash(candidate_nonce, siloed_note_hash);
+
+            if unique_note_hash == expected_unique_note_hash {
+                // Note that while we did check that the note hash is the preimage of the expected unique note hash, we
+                // perform no validations on the nullifier - we fundamentally cannot, since only the application knows
+                // how to compute nullifiers. We simply trust it to have provided the correct one: if it hasn't, then
+                // PXE may fail to realize that a given note has been nullified already, and calls to the application
+                // could result in invalid transactions (with duplicate nullifiers). This is not a concern because an
+                // application already has more direct means of making a call to it fail the transaction.
+                discovered_notes.push(
+                    DiscoveredNoteInfo {
+                        nonce: candidate_nonce,
+                        note_hash: hashes.note_hash,
+                        inner_nullifier: hashes.inner_nullifier,
+                    },
+                );
+
+                // We don't exit the loop - it is possible (though rare) for the exact same note content to be present
+                // multiple times in the same transaction with different nonces. This typically doesn't happen due to
+                // notes containing random values in order to hide their contents.
+            }
+        },
+    );
+
+    debug_log_format(
+        "Discovered a total of {0} notes",
+        [discovered_notes.len() as Field],
+    );
+
+    *discovered_notes
+}

aztec/src/discovery/partial_notes.nr

@@ -0,0 +1,153 @@
+use crate::{
+    capsules::CapsuleArray,
+    discovery::{
+        ComputeNoteHashAndNullifier,
+        nonce_discovery::{attempt_note_nonce_discovery, DiscoveredNoteInfo},
+        private_logs::MAX_PARTIAL_NOTE_PRIVATE_PACKED_LEN,
+    },
+    oracle::note_discovery::{deliver_note, get_log_by_tag},
+    utils::array,
+};
+
+use dep::protocol_types::{
+    address::AztecAddress,
+    constants::PUBLIC_LOG_DATA_SIZE_IN_FIELDS,
+    debug_log::debug_log_format,
+    traits::{Deserialize, Serialize, ToField},
+};
+
+/// The slot in the PXE capsules where we store a `CapsuleArray` of `DeliveredPendingPartialNote`.
+// TODO(#11630): come up with some sort of slot allocation scheme.
+pub global DELIVERED_PENDING_PARTIAL_NOTE_ARRAY_LENGTH_CAPSULES_SLOT: Field = 77;
+
+/// Public logs contain an extra field at the beginning with the address of the contract that emitted them, and partial
+/// notes emit their completion tag in the log, resulting in the first two fields in the public log not being part of
+/// the packed public content.
+// TODO(#10273): improve how contract log siloing is handled
+pub global NON_PACKED_CONTENT_FIELDS_IN_PUBLIC_LOG: u32 = 2;
+
+/// The maximum length of the packed representation of public fields in a partial note. This is limited by public log
+/// size and extra fields in the log (e.g. the tag).
+pub global MAX_PUBLIC_PARTIAL_NOTE_PACKED_CONTENT_LENGTH: u32 =
+    PUBLIC_LOG_DATA_SIZE_IN_FIELDS - NON_PACKED_CONTENT_FIELDS_IN_PUBLIC_LOG;
+
+/// A partial note that was delivered but is still pending completion. Contains the information necessary to find the
+/// log that will complete it and lead to a note being discovered and delivered.
+#[derive(Serialize, Deserialize)]
+pub(crate) struct DeliveredPendingPartialNote {
+    pub(crate) note_completion_log_tag: Field,
+    pub(crate) storage_slot: Field,
+    pub(crate) note_type_id: Field,
+    pub(crate) packed_private_note_content: BoundedVec<Field, MAX_PARTIAL_NOTE_PRIVATE_PACKED_LEN>,
+    pub(crate) recipient: AztecAddress,
+}
+
+/// Searches for public logs that would result in the completion of pending partial notes, ultimately resulting in the
+/// notes being delivered to PXE if completed.
+pub unconstrained fn fetch_and_process_public_partial_note_completion_logs<Env>(
+    contract_address: AztecAddress,
+    compute_note_hash_and_nullifier: ComputeNoteHashAndNullifier<Env>,
+) {
+    let pending_partial_notes = CapsuleArray::at(
+        contract_address,
+        DELIVERED_PENDING_PARTIAL_NOTE_ARRAY_LENGTH_CAPSULES_SLOT,
+    );
+
+    debug_log_format(
+        "{} pending partial notes",
+        [pending_partial_notes.len() as Field],
+    );
+
+    let mut i = &mut 0;
+    whyle(
+        || *i < pending_partial_notes.len(),
+        || {
+            let pending_partial_note: DeliveredPendingPartialNote = pending_partial_notes.get(*i);
+
+            let maybe_log = get_log_by_tag(pending_partial_note.note_completion_log_tag);
+            if maybe_log.is_none() {
+                debug_log_format(
+                    "Found no completion logs for partial note #{}",
+                    [(*i) as Field],
+                );
+                *i += 1 as u32;
+                // Note that we're not removing the pending partial note from the PXE DB, so we will continue searching
+                // for this tagged log when performing note discovery in the future until we either find it or the entry
+                // is somehow removed from the PXE DB.
+            } else {
+                debug_log_format("Completion log found for partial note #{}", [(*i) as Field]);
+                let log = maybe_log.unwrap();
+
+                // Public logs have an extra field at the beginning with the contract address, which we use to verify
+                // that we're getting the logs from the expected contract.
+                // TODO(#10273): improve how contract log siloing is handled
+                assert_eq(
+                    log.log_content.get(0),
+                    contract_address.to_field(),
+                    "Got a public log emitted by a different contract",
+                );
+
+                // Public fields are assumed to all be placed at the end of the packed representation, so we combine the
+                // private and public packed fields (i.e. the contents of the log sans the extra fields) to get the
+                // complete packed content.
+                let packed_public_note_content: BoundedVec<_, MAX_PUBLIC_PARTIAL_NOTE_PACKED_CONTENT_LENGTH> =
+                    array::subbvec(log.log_content, NON_PACKED_CONTENT_FIELDS_IN_PUBLIC_LOG);
+                let complete_packed_note_content = array::append(
+                    pending_partial_note.packed_private_note_content,
+                    packed_public_note_content,
+                );
+
+                let discovered_notes = attempt_note_nonce_discovery(
+                    log.unique_note_hashes_in_tx,
+                    log.first_nullifier_in_tx,
+                    compute_note_hash_and_nullifier,
+                    contract_address,
+                    pending_partial_note.storage_slot,
+                    pending_partial_note.note_type_id,
+                    complete_packed_note_content,
+                );
+
+                debug_log_format(
+                    "Discovered {0} notes for partial note {1}",
+                    [discovered_notes.len() as Field, (*i) as Field],
+                );
+
+                array::for_each_in_bounded_vec(
+                    discovered_notes,
+                    |discovered_note: DiscoveredNoteInfo, _| {
+                        // TODO:(#10728): decide how to handle notes that fail delivery. This could be due to e.g. a
+                        // temporary node connectivity issue - is simply throwing good enough here?
+                        assert(
+                            deliver_note(
+                                contract_address,
+                                pending_partial_note.storage_slot,
+                                discovered_note.nonce,
+                                complete_packed_note_content,
+                                discovered_note.note_hash,
+                                discovered_note.inner_nullifier,
+                                log.tx_hash,
+                                pending_partial_note.recipient,
+                            ),
+                            "Failed to deliver note",
+                        );
+                    },
+                );
+
+                // Because there is only a single log for a given tag, once we've processed the tagged log then we
+                // simply delete the pending work entry, regardless of whether it was actually completed or not.
+                // TODO(#11627): only remove the pending entry if we actually process a log that results in the note
+                // being completed.
+                pending_partial_notes.remove(*i);
+            }
+        },
+    );
+}
+
+/// Custom version of a while loop, calls `body` repeatedly until `condition` returns false. To be removed once Noir
+/// supports looping in unconstrained code.
+fn whyle<Env, Env2>(condition: fn[Env]() -> bool, body: fn[Env2]() -> ()) {
+    if condition() {
+        body();
+        whyle(condition, body);
+    }
+}