Addressed review comments for #390 (#396)

* fix(WIP): review comments

* fix(WIP): review comments

* docs: document under what circumstances SMT creating can fail

* fix: review comments

* fix: compare note hashes, return error on duplicate unauthenticated note IDs

* refactor: update batch creation logic

* refactor: migrate to the latest `miden-base`

* refactor: store input notes commitment in transaction batch

* fix: review comments

* refactor: address review comments

crates/block-producer/src/batch_builder/batch.rs

@@ -1,12 +1,18 @@
-use std::collections::BTreeSet;
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    mem,
+};
 
 use miden_objects::{
     accounts::AccountId,
     batches::BatchNoteTree,
     block::BlockAccountUpdate,
-    crypto::hash::blake::{Blake3Digest, Blake3_256},
-    notes::{NoteId, Nullifier},
-    transaction::{OutputNote, TransactionId, TxAccountUpdate},
+    crypto::{
+        hash::blake::{Blake3Digest, Blake3_256},
+        merkle::MerklePath,
+    },
+    notes::{NoteHeader, NoteId, Nullifier},
+    transaction::{InputNoteCommitment, OutputNote, TransactionId, TxAccountUpdate},
     Digest, MAX_NOTES_PER_BATCH,
 };
 use tracing::instrument;
@@ -26,80 +32,100 @@ pub type BatchId = Blake3Digest<32>;
 pub struct TransactionBatch {
     id: BatchId,
     updated_accounts: Vec<(TransactionId, TxAccountUpdate)>,
-    unauthenticated_input_notes: BTreeSet<NoteId>,
-    produced_nullifiers: Vec<Nullifier>,
+    input_notes: Vec<InputNoteCommitment>,
     output_notes_smt: BatchNoteTree,
     output_notes: Vec<OutputNote>,
 }
 
 impl TransactionBatch {
-    // CONSTRUCTOR
+    // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
+
     /// Returns a new [TransactionBatch] instantiated from the provided vector of proven
-    /// transactions.
+    /// transactions. If a map of unauthenticated notes found in the store is provided, it is used
+    /// for transforming unauthenticated notes into authenticated notes.
     ///
     /// # Errors
     /// Returns an error if:
-    /// - The number of created notes across all transactions exceeds 4096.
+    /// - The number of output notes across all transactions exceeds 4096.
+    /// - There are duplicated output notes or unauthenticated notes found across all transactions
+    ///   in the batch.
+    /// - Hashes for corresponding input notes and output notes don't match.
     ///
     /// TODO: enforce limit on the number of created nullifiers.
     #[instrument(target = "miden-block-producer", name = "new_batch", skip_all, err)]
-    pub fn new(txs: Vec<ProvenTransaction>) -> Result<Self, BuildBatchError> {
+    pub fn new(
+        txs: Vec<ProvenTransaction>,
+        found_unauthenticated_notes: Option<BTreeMap<NoteId, MerklePath>>,
+    ) -> Result<Self, BuildBatchError> {
         let id = Self::compute_id(&txs);
 
+        // Populate batch output notes and updated accounts.
+        let mut output_notes = OutputNoteTracker::new(&txs)?;
         let mut updated_accounts = vec![];
-        let mut produced_nullifiers = vec![];
         let mut unauthenticated_input_notes = BTreeSet::new();
         for tx in &txs {
             // TODO: we need to handle a possibility that a batch contains multiple transactions against
             //       the same account (e.g., transaction `x` takes account from state `A` to `B` and
             //       transaction `y` takes account from state `B` to `C`). These will need to be merged
             //       into a single "update" `A` to `C`.
             updated_accounts.push((tx.id(), tx.account_update().clone()));
-
-            for note in tx.input_notes() {
-                produced_nullifiers.push(note.nullifier());
-                if let Some(header) = note.header() {
-                    if !unauthenticated_input_notes.insert(header.id()) {
-                        return Err(BuildBatchError::DuplicateUnauthenticatedNote(
-                            header.id(),
-                            txs,
-                        ));
-                    }
+            // Check unauthenticated input notes for duplicates:
+            for note in tx.get_unauthenticated_notes() {
+                let id = note.id();
+                if !unauthenticated_input_notes.insert(id) {
+                    return Err(BuildBatchError::DuplicateUnauthenticatedNote(id, txs.clone()));
                 }
             }
         }
 
-        // Populate batch output notes, filtering out unauthenticated notes consumed in the same batch.
-        // Consumed notes are also removed from the unauthenticated input notes set in order to avoid
-        // consumption of notes with the same ID by one single input.
+        // Populate batch produced nullifiers and match output notes with corresponding
+        // unauthenticated input notes in the same batch, which are removed from the unauthenticated
+        // input notes set. We also don't add nullifiers for such output notes to the produced
+        // nullifiers set.
         //
         // One thing to note:
         // This still allows transaction `A` to consume an unauthenticated note `x` and output note `y`
         // and for transaction `B` to consume an unauthenticated note `y` and output note `x`
         // (i.e., have a circular dependency between transactions), but this is not a problem.
-        let output_notes: Vec<_> = txs
-            .iter()
-            .flat_map(|tx| tx.output_notes().iter())
-            .filter(|&note| !unauthenticated_input_notes.remove(&note.id()))
-            .cloned()
-            .collect();
+        let mut input_notes = vec![];
+        for input_note in txs.iter().flat_map(|tx| tx.input_notes().iter()) {
+            // Header is presented only for unauthenticated input notes.
+            let input_note = match input_note.header() {
+                Some(input_note_header) => {
+                    if output_notes.remove_note(input_note_header, &txs)? {
+                        continue;
+                    }
+
+                    match found_unauthenticated_notes {
+                        Some(ref found_notes) => match found_notes.get(&input_note_header.id()) {
+                            Some(_path) => input_note.nullifier().into(),
+                            None => input_note.clone(),
+                        },
+                        None => input_note.clone(),
+                    }
+                },
+                None => input_note.clone(),
+            };
+            input_notes.push(input_note)
+        }
+
+        let output_notes = output_notes.into_notes();
 
         if output_notes.len() > MAX_NOTES_PER_BATCH {
             return Err(BuildBatchError::TooManyNotesCreated(output_notes.len(), txs));
         }
 
-        // TODO: document under what circumstances SMT creating can fail
+        // Build the output notes SMT.
         let output_notes_smt = BatchNoteTree::with_contiguous_leaves(
             output_notes.iter().map(|note| (note.id(), note.metadata())),
         )
-        .map_err(|e| BuildBatchError::NotesSmtError(e, txs))?;
+        .expect("Unreachable: fails only if the output note list contains duplicates");
 
         Ok(Self {
             id,
             updated_accounts,
-            unauthenticated_input_notes,
-            produced_nullifiers,
+            input_notes,
             output_notes_smt,
             output_notes,
         })
@@ -134,14 +160,15 @@ impl TransactionBatch {
         })
     }
 
-    /// Returns unauthenticated input notes set consumed by the transactions in this batch.
-    pub fn unauthenticated_input_notes(&self) -> &BTreeSet<NoteId> {
-        &self.unauthenticated_input_notes
+    /// Returns input notes list consumed by the transactions in this batch. Any unauthenticated
+    /// input notes which have matching output notes within this batch are not included in this list.
+    pub fn input_notes(&self) -> &[InputNoteCommitment] {
+        &self.input_notes
     }
 
     /// Returns an iterator over produced nullifiers for all consumed notes.
     pub fn produced_nullifiers(&self) -> impl Iterator<Item = Nullifier> + '_ {
-        self.produced_nullifiers.iter().cloned()
+        self.input_notes.iter().map(InputNoteCommitment::nullifier)
     }
 
     /// Returns the root hash of the output notes SMT.
@@ -165,3 +192,55 @@ impl TransactionBatch {
         Blake3_256::hash(&buf)
     }
 }
+
+struct OutputNoteTracker {
+    output_notes: Vec<Option<OutputNote>>,
+    output_note_index: BTreeMap<NoteId, usize>,
+}
+
+impl OutputNoteTracker {
+    fn new(txs: &[ProvenTransaction]) -> Result<Self, BuildBatchError> {
+        let mut output_notes = vec![];
+        let mut output_note_index = BTreeMap::new();
+        for tx in txs {
+            for note in tx.output_notes().iter() {
+                if output_note_index.insert(note.id(), output_notes.len()).is_some() {
+                    return Err(BuildBatchError::DuplicateOutputNote(note.id(), txs.to_vec()));
+                }
+                output_notes.push(Some(note.clone()));
+            }
+        }
+
+        Ok(Self { output_notes, output_note_index })
+    }
+
+    pub fn remove_note(
+        &mut self,
+        input_note_header: &NoteHeader,
+        txs: &[ProvenTransaction],
+    ) -> Result<bool, BuildBatchError> {
+        let id = input_note_header.id();
+        if let Some(note_index) = self.output_note_index.remove(&id) {
+            if let Some(output_note) = mem::take(&mut self.output_notes[note_index]) {
+                let input_hash = input_note_header.hash();
+                let output_hash = output_note.hash();
+                if output_hash != input_hash {
+                    return Err(BuildBatchError::NoteHashesMismatch {
+                        id,
+                        input_hash,
+                        output_hash,
+                        txs: txs.to_vec(),
+                    });
+                }
+
+                return Ok(true);
+            }
+        }
+
+        Ok(false)
+    }
+
+    pub fn into_notes(self) -> Vec<OutputNote> {
+        self.output_notes.into_iter().flatten().collect()
+    }
+}

crates/block-producer/src/batch_builder/mod.rs

@@ -169,23 +169,28 @@ where
         // TODO: this can be optimized by first computing dangling notes of the batch itself,
         //       and only then checking against the other ready batches
         let dangling_notes = self.find_dangling_notes(&txs).await;
-        if !dangling_notes.is_empty() {
-            let stored_notes =
-                match self.store.get_note_authentication_info(dangling_notes.iter()).await {
-                    Ok(stored_notes) => stored_notes,
-                    Err(err) => return Err(BuildBatchError::NotePathsError(err, txs)),
-                };
-            let missing_notes: Vec<_> = dangling_notes
-                .into_iter()
-                .filter(|note_id| !stored_notes.contains_key(note_id))
-                .collect();
-
-            if !missing_notes.is_empty() {
-                return Err(BuildBatchError::UnauthenticatedNotesNotFound(missing_notes, txs));
-            }
-        }
+        let found_unauthenticated_notes = match dangling_notes.is_empty() {
+            true => None,
+            false => {
+                let stored_notes =
+                    match self.store.get_note_authentication_info(dangling_notes.iter()).await {
+                        Ok(stored_notes) => stored_notes,
+                        Err(err) => return Err(BuildBatchError::NotePathsError(err, txs)),
+                    };
+                let missing_notes: Vec<_> = dangling_notes
+                    .into_iter()
+                    .filter(|note_id| !stored_notes.contains_key(note_id))
+                    .collect();
+
+                if !missing_notes.is_empty() {
+                    return Err(BuildBatchError::UnauthenticatedNotesNotFound(missing_notes, txs));
+                }
+
+                Some(stored_notes)
+            },
+        };
 
-        let batch = TransactionBatch::new(txs)?;
+        let batch = TransactionBatch::new(txs, found_unauthenticated_notes)?;
 
         info!(target: COMPONENT, "Transaction batch built");
         Span::current().record("batch_id", format_blake3_digest(batch.id()));

crates/block-producer/src/batch_builder/tests/mod.rs

@@ -155,5 +155,5 @@ fn dummy_tx_batch(starting_account_index: u32, num_txs_in_batch: usize) -> Trans
             MockProvenTxBuilder::with_account_index(starting_account_index + index as u32).build()
         })
         .collect();
-    TransactionBatch::new(txs).unwrap()
+    TransactionBatch::new(txs, None).unwrap()
 }

crates/block-producer/src/block_builder/mod.rs

@@ -4,7 +4,8 @@ use async_trait::async_trait;
 use miden_node_utils::formatting::{format_array, format_blake3_digest};
 use miden_objects::{
     block::{Block, BlockAccountUpdate},
-    notes::Nullifier,
+    notes::{NoteHeader, Nullifier},
+    transaction::InputNoteCommitment,
 };
 use tracing::{debug, info, instrument};
 
@@ -76,24 +77,29 @@ where
         let updated_accounts: Vec<_> =
             batches.iter().flat_map(TransactionBatch::updated_accounts).collect();
 
-        let created_notes: Vec<_> =
+        let output_notes: Vec<_> =
             batches.iter().map(TransactionBatch::output_notes).cloned().collect();
 
         let produced_nullifiers: Vec<Nullifier> =
             batches.iter().flat_map(TransactionBatch::produced_nullifiers).collect();
 
-        let created_notes_set: BTreeSet<_> = created_notes
+        // Populate set of output notes from all batches
+        let output_notes_set: BTreeSet<_> = output_notes
             .iter()
             .flat_map(|batch| batch.iter().map(|note| note.id()))
             .collect();
 
+        // Build a set of unauthenticated input notes for this block which do not have a matching
+        // output note produced in this block
         let dangling_notes: BTreeSet<_> = batches
             .iter()
-            .flat_map(TransactionBatch::unauthenticated_input_notes)
-            .filter(|&note_id| !created_notes_set.contains(note_id))
-            .copied()
+            .flat_map(TransactionBatch::input_notes)
+            .filter_map(InputNoteCommitment::header)
+            .map(NoteHeader::id)
+            .filter(|note_id| !output_notes_set.contains(note_id))
             .collect();
 
+        // Request information needed for block building from the store
         let block_inputs = self
             .store
             .get_block_inputs(
@@ -119,7 +125,7 @@ where
 
         // TODO: return an error?
         let block =
-            Block::new(new_block_header, updated_accounts, created_notes, produced_nullifiers)
+            Block::new(new_block_header, updated_accounts, output_notes, produced_nullifiers)
                 .expect("invalid block components");
 
         let block_hash = block.hash();