feat(block-producer): batch builder

crates/block-producer/Cargo.toml

@@ -41,6 +41,7 @@ miden-lib = { workspace = true, features = ["testing"] }
 miden-node-test-macro = { path = "../test-macro" }
 miden-objects = { workspace = true, features = ["testing"] }
 miden-tx = { workspace = true, features = ["testing"] }
+rand = "0.8.5"
 rand_chacha = { version = "0.3", default-features = false }
 tokio = { workspace = true, features = ["test-util"] }
 winterfell = { version = "0.9" }

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

@@ -22,7 +22,8 @@ pub type BatchId = Blake3Digest<32>;
 // TRANSACTION BATCH
 // ================================================================================================
 
-/// A batch of transactions that share a common proof.
+/// A batch of transactions that share a common proof. For any given account, at most 1 transaction
+/// in the batch must be addressing that account (issue: #186).
 ///
 /// Note: Until recursive proofs are available in the Miden VM, we don't include the common proof.
 #[derive(Debug, Clone, PartialEq, Eq)]

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

@@ -1,10 +1,18 @@
-use std::{cmp::min, collections::BTreeSet, num::NonZeroUsize, sync::Arc, time::Duration};
+use std::{
+    cmp::min,
+    collections::{BTreeMap, BTreeSet},
+    num::NonZeroUsize,
+    sync::Arc,
+    time::Duration,
+};
 
 use miden_objects::{
+    accounts::AccountId,
     notes::NoteId,
     transaction::{OutputNote, TransactionId},
+    Digest,
 };
-use tokio::{sync::Mutex, time};
+use tokio::{sync::Mutex, task::JoinSet, time};
 use tonic::async_trait;
 use tracing::{debug, info, instrument, Span};
 
@@ -219,31 +227,53 @@ pub struct BatchProducer {
     pub workers: NonZeroUsize,
     pub mempool: Arc<Mutex<Mempool>>,
     pub tx_per_batch: usize,
+    pub simulated_proof_time: Duration,
 }
 
-type BatchResult = Result<BatchJobId, (BatchJobId, BuildBatchError)>;
+type BatchResult = Result<(BatchJobId, TransactionBatch), (BatchJobId, BuildBatchError)>;
 
 /// Wrapper around tokio's JoinSet that remains pending if the set is empty,
 /// instead of returning None.
-struct WorkerPool(tokio::task::JoinSet<BatchResult>);
+struct WorkerPool {
+    in_progress: JoinSet<BatchResult>,
+    simulated_proof_time: Duration,
+}
 
 impl WorkerPool {
+    fn new(simulated_proof_time: Duration) -> Self {
+        Self {
+            simulated_proof_time,
+            in_progress: JoinSet::new(),
+        }
+    }
+
     async fn join_next(&mut self) -> Result<BatchResult, tokio::task::JoinError> {
-        if self.0.is_empty() {
+        if self.in_progress.is_empty() {
             std::future::pending().await
         } else {
             // Cannot be None as its not empty.
-            self.0.join_next().await.unwrap()
+            self.in_progress.join_next().await.unwrap()
         }
     }
 
     fn len(&self) -> usize {
-        self.0.len()
+        self.in_progress.len()
     }
 
-    fn spawn(&mut self, id: BatchJobId, transactions: Vec<TransactionId>) {
-        self.0.spawn(async move {
-            todo!("Do actual work like aggregating transaction data");
+    fn spawn(&mut self, id: BatchJobId, transactions: Vec<ProvenTransaction>) {
+        self.in_progress.spawn({
+            let simulated_proof_time = self.simulated_proof_time;
+            async move {
+                tracing::debug!("Begin proving batch.");
+
+                let batch = TransactionBatch::new(transactions, Default::default())
+                    .map_err(|err| (id, err))?;
+
+                tokio::time::sleep(simulated_proof_time).await;
+                tracing::debug!("Batch proof completed.");
+
+                Ok((id, batch))
+            }
         });
     }
 }
@@ -253,7 +283,7 @@ impl BatchProducer {
         let mut interval = tokio::time::interval(self.batch_interval);
         interval.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
 
-        let mut inflight = WorkerPool(tokio::task::JoinSet::new());
+        let mut inflight = WorkerPool::new(self.simulated_proof_time);
 
         loop {
             tokio::select! {
@@ -285,8 +315,8 @@ impl BatchProducer {
                             tracing::warn!(%batch_id, %err, "Batch job failed.");
                             mempool.batch_failed(batch_id);
                         },
-                        Ok(Ok(batch_id)) => {
-                            mempool.batch_proved(batch_id);
+                        Ok(Ok((batch_id, batch))) => {
+                            mempool.batch_proved(batch_id, batch);
                         }
                     }
                 }

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

@@ -1,4 +1,7 @@
-use std::{collections::BTreeSet, sync::Arc};
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    sync::Arc,
+};
 
 use async_trait::async_trait;
 use miden_node_utils::formatting::{format_array, format_blake3_digest};
@@ -146,12 +149,13 @@ where
     }
 }
 
-struct BlockProducer {
+struct BlockProducer<BB> {
     pub mempool: Arc<Mutex<Mempool>>,
     pub block_interval: tokio::time::Duration,
+    pub block_builder: BB,
 }
 
-impl BlockProducer {
+impl<BB: BlockBuilder> BlockProducer<BB> {
     pub async fn run(self) {
         let mut interval = tokio::time::interval(self.block_interval);
         interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
@@ -160,8 +164,9 @@ impl BlockProducer {
             interval.tick().await;
 
             let (block_number, batches) = self.mempool.lock().await.select_block();
+            let batches = batches.into_values().collect::<Vec<_>>();
 
-            let result = self.build_and_commit_block(batches).await;
+            let result = self.block_builder.build_block(&batches).await;
             let mut mempool = self.mempool.lock().await;
 
             match result {
@@ -170,8 +175,4 @@ impl BlockProducer {
             }
         }
     }
-
-    async fn build_and_commit_block(&self, batches: BTreeSet<BatchJobId>) -> Result<(), ()> {
-        todo!("Aggregate, prove and commit block");
-    }
 }

crates/block-producer/src/errors.rs

@@ -1,5 +1,3 @@
-use std::collections::BTreeSet;
-
 use miden_node_proto::errors::ConversionError;
 use miden_node_utils::formatting::format_opt;
 use miden_objects::{
@@ -11,7 +9,6 @@ use miden_objects::{
     MAX_BATCHES_PER_BLOCK, MAX_INPUT_NOTES_PER_BATCH, MAX_OUTPUT_NOTES_PER_BATCH,
 };
 use miden_processor::ExecutionError;
-use miden_tx::TransactionVerifierError;
 use thiserror::Error;
 
 use crate::mempool::BlockNumber;
@@ -32,6 +29,9 @@ pub enum VerifyTxError {
     )]
     UnauthenticatedNotesNotFound(Vec<NoteId>),
 
+    #[error("Output note IDs already used: {0:?}")]
+    OutputNotesAlreadyExist(Vec<NoteId>),
+
     /// The account's initial hash did not match the current account's hash
     #[error("Incorrect account's initial hash ({tx_initial_account_hash}, current: {})", format_opt(.current_account_hash.as_ref()))]
     IncorrectAccountInitialHash {
@@ -61,31 +61,36 @@ pub enum VerifyTxError {
 pub enum AddTransactionError {
     #[error("Transaction verification failed: {0}")]
     VerificationFailed(#[from] VerifyTxError),
-}
 
-#[derive(thiserror::Error, Debug, PartialEq)]
-pub enum AddTransactionErrorRework {
-    #[error("Transaction's initial account state {expected} did not match the current account state {current}.")]
-    InvalidAccountState { current: Digest, expected: Digest },
     #[error("Transaction input data is stale. Required data fresher than {stale_limit} but inputs are from {input_block}.")]
     StaleInputs {
         input_block: BlockNumber,
         stale_limit: BlockNumber,
     },
-    #[error("Authenticated note nullifier {0} not found.")]
-    AuthenticatedNoteNotFound(Nullifier),
-    #[error("Unauthenticated note {0} not found.")]
-    UnauthenticatedNoteNotFound(NoteId),
-    #[error("Note nullifiers already consumed: {0:?}")]
-    NotesAlreadyConsumed(BTreeSet<Nullifier>),
-    #[error(transparent)]
-    TxInputsError(#[from] TxInputsError),
-    #[error(transparent)]
-    ProofVerificationFailed(#[from] TransactionVerifierError),
-    #[error("Failed to deserialize transaction: {0}.")]
+
+    #[error("Deserialization failed: {0}")]
     DeserializationError(String),
 }
 
+impl From<AddTransactionError> for tonic::Status {
+    fn from(value: AddTransactionError) -> Self {
+        use AddTransactionError::*;
+        match value {
+            VerificationFailed(VerifyTxError::InputNotesAlreadyConsumed(_))
+            | VerificationFailed(VerifyTxError::UnauthenticatedNotesNotFound(_))
+            | VerificationFailed(VerifyTxError::OutputNotesAlreadyExist(_))
+            | VerificationFailed(VerifyTxError::IncorrectAccountInitialHash { .. })
+            | VerificationFailed(VerifyTxError::InvalidTransactionProof(_))
+            | DeserializationError(_) => Self::invalid_argument(value.to_string()),
+
+            // Internal errors which should not be communicated to the user.
+            VerificationFailed(VerifyTxError::TransactionInputError(_))
+            | VerificationFailed(VerifyTxError::StoreConnectionFailed(_))
+            | StaleInputs { .. } => Self::internal("Internal error"),
+        }
+    }
+}
+
 // Batch building errors
 // =================================================================================================
 

crates/block-producer/src/mempool/batch_graph.rs

@@ -4,6 +4,10 @@ use miden_objects::transaction::TransactionId;
 use miden_tx::utils::collections::KvMap;
 
 use super::BatchJobId;
+use crate::batch_builder::batch::TransactionBatch;
+
+// BATCH GRAPH
+// ================================================================================================
 
 #[derive(Default, Clone)]
 pub struct BatchGraph {
@@ -53,12 +57,12 @@ impl BatchGraph {
 
         // New node might be a root.
         //
-        // This could be optimised by inlining this inside the parent loop. This would prevent the
+        // This could be optimized by inlining this inside the parent loop. This would prevent the
         // double iteration over parents, at the cost of some code duplication.
         self.try_make_root(id);
     }
 
-    /// Removes the batches and all their descendents from the graph.
+    /// Removes the batches and all their descendants from the graph.
     ///
     /// Returns all removed batches and their transactions.
     pub fn purge_subgraphs(
@@ -76,8 +80,8 @@ impl BatchGraph {
                 continue;
             };
 
-            // All the child batches are also removed so no need to chec
-            // for new roots. No new roots are possible as a result of this subgraph removal.
+            // All the child batches are also removed so no need to check for new roots. No new
+            // roots are possible as a result of this subgraph removal.
             self.roots.remove(&node_id);
 
             for transaction in &node.transactions {
@@ -86,8 +90,8 @@ impl BatchGraph {
 
             // Inform parent that this child no longer exists.
             //
-            // The same is not required for children of this batch as we will
-            // be removing those as well.
+            // The same is not required for children of this batch as we will be removing those as
+            // well.
             for parent in &node.parents {
                 // Parent could already be removed as part of this subgraph removal.
                 if let Some(parent) = self.nodes.get_mut(parent) {
@@ -102,7 +106,7 @@ impl BatchGraph {
         removed
     }
 
-    /// Removes a set of batches from the graph without removing any descendents.
+    /// Removes a set of batches from the graph without removing any descendants.
     ///
     /// This is intended to cull completed batches from stale blocs.
     pub fn remove_committed(&mut self, batches: BTreeSet<BatchJobId>) -> Vec<TransactionId> {
@@ -129,32 +133,36 @@ impl BatchGraph {
     }
 
     /// Mark a batch as proven if it exists.
-    pub fn mark_proven(&mut self, id: BatchJobId) {
-        // Its possible for inflight batches to have been removed as part
-        // of another batches failure.
+    pub fn mark_proven(&mut self, id: BatchJobId, batch: TransactionBatch) {
+        // Its possible for inflight batches to have been removed as part of another batches
+        // failure.
         if let Some(node) = self.nodes.get_mut(&id) {
-            node.status = Status::Proven;
+            node.status = Status::Proven(batch);
             self.try_make_root(id);
         }
     }
 
     /// Returns at most `count` __indepedent__ batches which are ready for inclusion in a block.
-    pub fn select_block(&mut self, count: usize) -> BTreeSet<BatchJobId> {
-        let mut batches = BTreeSet::new();
+    pub fn select_block(&mut self, count: usize) -> BTreeMap<BatchJobId, TransactionBatch> {
+        let mut batches = BTreeMap::new();
 
         // Track children so we can evaluate them for root afterwards.
         let mut children = BTreeSet::new();
 
-        for batch in &self.roots {
-            let mut node = self.nodes.get_mut(batch).expect("Root node must be in graph");
+        for batch_id in &self.roots {
+            let mut node = self.nodes.get_mut(batch_id).expect("Root node must be in graph");
 
             // Filter out batches which have dependencies in our selection so far.
-            if batches.union(&node.parents).next().is_some() {
+            if node.parents.iter().any(|parent| batches.contains_key(parent)) {
                 continue;
             }
 
-            batches.insert(*batch);
-            node.status = Status::Proven;
+            let Status::Proven(batch) = node.status.clone() else {
+                unreachable!("Root batch must be in proven state.");
+            };
+
+            batches.insert(*batch_id, batch);
+            node.status = Status::InBlock;
 
             if batches.len() == count {
                 break;
@@ -193,9 +201,9 @@ struct Node {
     children: BTreeSet<BatchJobId>,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 enum Status {
     InFlight,
-    Proven,
+    Proven(TransactionBatch),
     InBlock,
 }