fix: review comments

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

@@ -1,10 +1,16 @@
-use std::collections::{BTreeMap, BTreeSet};
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    mem,
+};
 
 use miden_objects::{
     accounts::AccountId,
     batches::BatchNoteTree,
     block::BlockAccountUpdate,
-    crypto::hash::blake::{Blake3Digest, Blake3_256},
+    crypto::{
+        hash::blake::{Blake3Digest, Blake3_256},
+        merkle::MerklePath,
+    },
     notes::{NoteId, Nullifier},
     transaction::{InputNoteCommitment, OutputNote, TransactionId, TxAccountUpdate},
     Digest, MAX_NOTES_PER_BATCH,
@@ -28,27 +34,36 @@ pub struct TransactionBatch {
     updated_accounts: Vec<(TransactionId, TxAccountUpdate)>,
     input_notes: Vec<InputNoteCommitment>,
     output_notes_smt: BatchNoteTree,
-    output_notes: BTreeMap<NoteId, OutputNote>,
+    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 updated_accounts = vec![];
-        let mut output_notes = BTreeMap::new();
+        let mut output_notes = vec![];
+        let mut output_note_index = BTreeMap::new();
         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
@@ -57,9 +72,10 @@ impl TransactionBatch {
             //       into a single "update" `A` to `C`.
             updated_accounts.push((tx.id(), tx.account_update().clone()));
             for note in tx.output_notes().iter() {
-                if output_notes.insert(note.id(), note.clone()).is_some() {
+                if output_note_index.insert(note.id(), output_notes.len()).is_some() {
                     return Err(BuildBatchError::DuplicateOutputNote(note.id(), txs.clone()));
                 }
+                output_notes.push(Some(note.clone()));
             }
             // Check unauthenticated input notes for duplicates:
             for note in tx.get_unauthenticated_notes() {
@@ -82,35 +98,49 @@ impl TransactionBatch {
         let mut input_notes = vec![];
         for input_note in txs.iter().flat_map(|tx| tx.input_notes().iter()) {
             // Header is presented only for unauthenticated notes.
-            if let Some(input_note_header) = input_note.header() {
-                let id = input_note_header.id();
-                if let Some(output_note) = output_notes.remove(&id) {
-                    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.clone(),
-                        });
+            let input_note = match input_note.header() {
+                Some(input_note_header) => {
+                    let id = input_note_header.id();
+                    if let Some(note_index) = output_note_index.remove(&id) {
+                        if let Some(output_note) = mem::take(&mut 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.clone(),
+                                });
+                            }
+
+                            // Don't add input notes if corresponding output notes consumed in the same batch.
+                            continue;
+                        }
                     }
 
-                    // Don't add input notes if corresponding output notes consumed in the same batch.
-                    continue;
-                }
-            }
-
-            input_notes.push(input_note.clone());
+                    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: Vec<_> = output_notes.into_iter().flatten().collect();
+
         if output_notes.len() > MAX_NOTES_PER_BATCH {
             return Err(BuildBatchError::TooManyNotesCreated(output_notes.len(), txs));
         }
 
         // Build the output notes SMT.
         let output_notes_smt = BatchNoteTree::with_contiguous_leaves(
-            output_notes.iter().map(|(&id, note)| (id, note.metadata())),
+            output_notes.iter().map(|note| (note.id(), note.metadata())),
         )
         .expect("Unreachable: fails only if the output note list contains duplicates");
 
@@ -168,7 +198,7 @@ impl TransactionBatch {
     }
 
     /// Returns output notes list.
-    pub fn output_notes(&self) -> &BTreeMap<NoteId, OutputNote> {
+    pub fn output_notes(&self) -> &Vec<OutputNote> {
         &self.output_notes
     }
 

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

@@ -132,7 +132,7 @@ where
                     .read()
                     .await
                     .iter()
-                    .flat_map(|batch| batch.output_notes().keys().copied()),
+                    .flat_map(|batch| batch.output_notes().iter().map(OutputNote::id)),
             )
             .collect();
 
@@ -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

@@ -3,7 +3,7 @@ use std::{collections::BTreeSet, sync::Arc};
 use async_trait::async_trait;
 use miden_node_utils::formatting::{format_array, format_blake3_digest};
 use miden_objects::{
-    block::{Block, BlockAccountUpdate, NoteBatch},
+    block::{Block, BlockAccountUpdate},
     notes::{NoteHeader, Nullifier},
     transaction::InputNoteCommitment,
 };
@@ -77,27 +77,29 @@ where
         let updated_accounts: Vec<_> =
             batches.iter().flat_map(TransactionBatch::updated_accounts).collect();
 
-        let created_notes: Vec<NoteBatch> = batches
-            .iter()
-            .map(|batch| batch.output_notes().values().cloned().collect())
-            .collect();
+        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::input_notes)
             .filter_map(InputNoteCommitment::header)
             .map(NoteHeader::id)
-            .filter(|note_id| !created_notes_set.contains(note_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(
@@ -123,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();

crates/block-producer/src/block_builder/prover/tests.rs

@@ -69,7 +69,7 @@ fn test_block_witness_validation_inconsistent_account_ids() {
             )
             .build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         let batch_2 = {
@@ -80,7 +80,7 @@ fn test_block_witness_validation_inconsistent_account_ids() {
             )
             .build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         vec![batch_1, batch_2]
@@ -134,19 +134,25 @@ fn test_block_witness_validation_inconsistent_account_hashes() {
     };
 
     let batches = {
-        let batch_1 = TransactionBatch::new(vec![MockProvenTxBuilder::with_account(
-            account_id_1,
-            account_1_hash_batches,
-            Digest::default(),
+        let batch_1 = TransactionBatch::new(
+            vec![MockProvenTxBuilder::with_account(
+                account_id_1,
+                account_1_hash_batches,
+                Digest::default(),
+            )
+            .build()],
+            None,
         )
-        .build()])
         .unwrap();
-        let batch_2 = TransactionBatch::new(vec![MockProvenTxBuilder::with_account(
-            account_id_2,
-            Digest::default(),
-            Digest::default(),
+        let batch_2 = TransactionBatch::new(
+            vec![MockProvenTxBuilder::with_account(
+                account_id_2,
+                Digest::default(),
+                Digest::default(),
+            )
+            .build()],
+            None,
         )
-        .build()])
         .unwrap();
 
         vec![batch_1, batch_2]
@@ -229,8 +235,8 @@ async fn test_compute_account_root_success() {
             })
             .collect();
 
-        let batch_1 = TransactionBatch::new(txs[..2].to_vec()).unwrap();
-        let batch_2 = TransactionBatch::new(txs[2..].to_vec()).unwrap();
+        let batch_1 = TransactionBatch::new(txs[..2].to_vec(), None).unwrap();
+        let batch_2 = TransactionBatch::new(txs[2..].to_vec(), None).unwrap();
 
         vec![batch_1, batch_2]
     };
@@ -373,7 +379,7 @@ async fn test_compute_note_root_empty_notes_success() {
         .unwrap();
 
     let batches: Vec<TransactionBatch> = {
-        let batch = TransactionBatch::new(Vec::new()).unwrap();
+        let batch = TransactionBatch::new(vec![], None).unwrap();
         vec![batch]
     };
 
@@ -446,8 +452,8 @@ async fn test_compute_note_root_success() {
             })
             .collect();
 
-        let batch_1 = TransactionBatch::new(txs[..2].to_vec()).unwrap();
-        let batch_2 = TransactionBatch::new(txs[2..].to_vec()).unwrap();
+        let batch_1 = TransactionBatch::new(txs[..2].to_vec(), None).unwrap();
+        let batch_2 = TransactionBatch::new(txs[2..].to_vec(), None).unwrap();
 
         vec![batch_1, batch_2]
     };
@@ -495,13 +501,13 @@ fn test_block_witness_validation_inconsistent_nullifiers() {
         let batch_1 = {
             let tx = MockProvenTxBuilder::with_account_index(0).nullifiers_range(0..1).build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         let batch_2 = {
             let tx = MockProvenTxBuilder::with_account_index(1).nullifiers_range(1..2).build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         vec![batch_1, batch_2]
@@ -571,13 +577,13 @@ async fn test_compute_nullifier_root_empty_success() {
         let batch_1 = {
             let tx = MockProvenTxBuilder::with_account_index(0).build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         let batch_2 = {
             let tx = MockProvenTxBuilder::with_account_index(1).build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         vec![batch_1, batch_2]
@@ -624,13 +630,13 @@ async fn test_compute_nullifier_root_success() {
         let batch_1 = {
             let tx = MockProvenTxBuilder::with_account_index(0).nullifiers_range(0..1).build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         let batch_2 = {
             let tx = MockProvenTxBuilder::with_account_index(1).nullifiers_range(1..2).build();
 
-            TransactionBatch::new(vec![tx]).unwrap()
+            TransactionBatch::new(vec![tx], None).unwrap()
         };
 
         vec![batch_1, batch_2]