block_builder: using borrows instead of refcount (0xPolygonMiden#188)

block-producer/src/batch_builder/batch.rs

@@ -6,7 +6,7 @@ use miden_vm::crypto::SimpleSmt;
 use tracing::instrument;
 
 use super::errors::BuildBatchError;
-use crate::{SharedProvenTx, CREATED_NOTES_SMT_DEPTH, MAX_NUM_CREATED_NOTES_PER_BATCH};
+use crate::{ProvenTransaction, CREATED_NOTES_SMT_DEPTH, MAX_NUM_CREATED_NOTES_PER_BATCH};
 
 pub type BatchId = Blake3Digest<32>;
 
@@ -17,7 +17,7 @@ pub type BatchId = Blake3Digest<32>;
 /// 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)]
+#[derive(Debug, Clone)]
 pub struct TransactionBatch {
     id: BatchId,
     updated_accounts: BTreeMap<AccountId, AccountStates>,
@@ -39,7 +39,9 @@ impl TransactionBatch {
     ///
     /// 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<SharedProvenTx>) -> Result<Self, BuildBatchError> {
+    pub fn new(txs: Vec<ProvenTransaction>) -> Result<Self, BuildBatchError> {
+        let id = Self::compute_id(&txs);
+
         let updated_accounts = txs
             .iter()
             .map(|tx| {
@@ -64,7 +66,7 @@ impl TransactionBatch {
                 txs.iter().flat_map(|tx| tx.output_notes().iter()).cloned().collect();
 
             if created_notes.len() > MAX_NUM_CREATED_NOTES_PER_BATCH {
-                return Err(BuildBatchError::TooManyNotesCreated(created_notes.len()));
+                return Err(BuildBatchError::TooManyNotesCreated(created_notes.len(), txs));
             }
 
             // TODO: document under what circumstances SMT creating can fail
@@ -74,12 +76,11 @@ impl TransactionBatch {
                     created_notes.into_iter().flat_map(|note_envelope| {
                         [note_envelope.note_id().into(), note_envelope.metadata().into()]
                     }),
-                )?,
+                )
+                .map_err(|e| BuildBatchError::NotesSmtError(e, txs))?,
             )
         };
 
-        let id = Self::compute_id(&txs);
-
         Ok(Self {
             id,
             updated_accounts,
@@ -131,7 +132,7 @@ impl TransactionBatch {
     // HELPER FUNCTIONS
     // --------------------------------------------------------------------------------------------
 
-    fn compute_id(txs: &[SharedProvenTx]) -> BatchId {
+    fn compute_id(txs: &[ProvenTransaction]) -> BatchId {
         let mut buf = Vec::with_capacity(32 * txs.len());
         for tx in txs {
             buf.extend_from_slice(&tx.id().as_bytes());
@@ -144,7 +145,7 @@ impl TransactionBatch {
 /// account.
 ///
 /// TODO: should this be moved into domain objects?
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 struct AccountStates {
     initial_state: Digest,
     final_state: Digest,

block-producer/src/batch_builder/errors.rs

@@ -1,15 +1,31 @@
+use miden_objects::transaction::ProvenTransaction;
 use miden_vm::crypto::MerkleError;
 use thiserror::Error;
 
 use crate::MAX_NUM_CREATED_NOTES_PER_BATCH;
 
-#[derive(Error, Debug, PartialEq)]
+/// Error that may happen while building a transaction batch.
+///
+/// These errors are returned from the batch builder to the transaction queue, instead of
+/// dropping the transactions, they are included into the error values, so that the transaction
+/// queue can re-queue them.
+#[derive(Error, Debug)]
 pub enum BuildBatchError {
     #[error(
         "Too many notes in the batch. Got: {0}, max: {}",
         MAX_NUM_CREATED_NOTES_PER_BATCH
     )]
-    TooManyNotesCreated(usize),
+    TooManyNotesCreated(usize, Vec<ProvenTransaction>),
+
     #[error("failed to create notes SMT: {0}")]
-    NotesSmtError(#[from] MerkleError),
+    NotesSmtError(MerkleError, Vec<ProvenTransaction>),
+}
+
+impl BuildBatchError {
+    pub fn into_transactions(self) -> Vec<ProvenTransaction> {
+        match self {
+            BuildBatchError::TooManyNotesCreated(_, txs) => txs,
+            BuildBatchError::NotesSmtError(_, txs) => txs,
+        }
+    }
 }

block-producer/src/batch_builder/mod.rs

@@ -5,13 +5,13 @@ use tokio::{sync::RwLock, time};
 use tracing::{debug, info, instrument, Span};
 
 use self::errors::BuildBatchError;
-use crate::{block_builder::BlockBuilder, SharedProvenTx, SharedRwVec, SharedTxBatch, COMPONENT};
+use crate::{block_builder::BlockBuilder, ProvenTransaction, SharedRwVec, COMPONENT};
 
 pub mod errors;
 #[cfg(test)]
 mod tests;
 
-mod batch;
+pub mod batch;
 pub use batch::TransactionBatch;
 use miden_node_utils::logging::{format_array, format_blake3_digest};
 
@@ -32,7 +32,7 @@ pub trait BatchBuilder: Send + Sync + 'static {
     /// Start proving of a new batch.
     async fn build_batch(
         &self,
-        txs: Vec<SharedProvenTx>,
+        txs: Vec<ProvenTransaction>,
     ) -> Result<(), BuildBatchError>;
 }
 
@@ -50,7 +50,7 @@ pub struct DefaultBatchBuilderOptions {
 
 pub struct DefaultBatchBuilder<BB> {
     /// Batches ready to be included in a block
-    ready_batches: SharedRwVec<SharedTxBatch>,
+    ready_batches: SharedRwVec<TransactionBatch>,
 
     block_builder: Arc<BB>,
 
@@ -97,7 +97,7 @@ where
     /// Note that we call `build_block()` regardless of whether the `ready_batches` queue is empty.
     /// A call to an empty `build_block()` indicates that an empty block should be created.
     async fn try_build_block(&self) {
-        let mut batches_in_block: Vec<SharedTxBatch> = {
+        let mut batches_in_block: Vec<TransactionBatch> = {
             let mut locked_ready_batches = self.ready_batches.write().await;
 
             let num_batches_in_block =
@@ -106,7 +106,7 @@ where
             locked_ready_batches.drain(..num_batches_in_block).collect()
         };
 
-        match self.block_builder.build_block(batches_in_block.clone()).await {
+        match self.block_builder.build_block(&batches_in_block).await {
             Ok(_) => {
                 // block successfully built, do nothing
             },
@@ -127,14 +127,14 @@ where
     #[instrument(target = "miden-block-producer", skip_all, err, fields(batch_id))]
     async fn build_batch(
         &self,
-        txs: Vec<SharedProvenTx>,
+        txs: Vec<ProvenTransaction>,
     ) -> Result<(), BuildBatchError> {
         let num_txs = txs.len();
 
         info!(target: COMPONENT, num_txs, "Building a transaction batch");
         debug!(target: COMPONENT, txs = %format_array(txs.iter().map(|tx| tx.id().to_hex())));
 
-        let batch = Arc::new(TransactionBatch::new(txs)?);
+        let batch = TransactionBatch::new(txs)?;
 
         info!(target: COMPONENT, "Transaction batch built");
         Span::current().record("batch_id", format_blake3_digest(batch.id()));

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

@@ -1,27 +1,27 @@
 use super::*;
 use crate::{
-    block_builder::errors::BuildBlockError, test_utils::DummyProvenTxGenerator, SharedTxBatch,
+    block_builder::errors::BuildBlockError, test_utils::DummyProvenTxGenerator, TransactionBatch,
 };
 
 // STRUCTS
 // ================================================================================================
 
 #[derive(Default)]
 struct BlockBuilderSuccess {
-    batch_groups: SharedRwVec<Vec<SharedTxBatch>>,
+    batch_groups: SharedRwVec<Vec<TransactionBatch>>,
     num_empty_batches_received: Arc<RwLock<usize>>,
 }
 
 #[async_trait]
 impl BlockBuilder for BlockBuilderSuccess {
     async fn build_block(
         &self,
-        batches: Vec<SharedTxBatch>,
+        batches: &[TransactionBatch],
     ) -> Result<(), BuildBlockError> {
         if batches.is_empty() {
             *self.num_empty_batches_received.write().await += 1;
         } else {
-            self.batch_groups.write().await.push(batches);
+            self.batch_groups.write().await.push(batches.to_vec());
         }
 
         Ok(())
@@ -35,7 +35,7 @@ struct BlockBuilderFailure;
 impl BlockBuilder for BlockBuilderFailure {
     async fn build_block(
         &self,
-        _batches: Vec<SharedTxBatch>,
+        _batches: &[TransactionBatch],
     ) -> Result<(), BuildBlockError> {
         Err(BuildBlockError::TooManyBatchesInBlock(0))
     }
@@ -163,8 +163,7 @@ async fn test_batches_added_back_to_queue_on_block_build_failure() {
 fn dummy_tx_batch(
     tx_gen: &DummyProvenTxGenerator,
     num_txs_in_batch: usize,
-) -> SharedTxBatch {
-    let txs: Vec<_> = (0..num_txs_in_batch).map(|_| Arc::new(tx_gen.dummy_proven_tx())).collect();
-
-    Arc::new(TransactionBatch::new(txs).unwrap())
+) -> TransactionBatch {
+    let txs: Vec<_> = (0..num_txs_in_batch).map(|_| tx_gen.dummy_proven_tx()).collect();
+    TransactionBatch::new(txs).unwrap()
 }

block-producer/src/block_builder/mod.rs

@@ -5,9 +5,10 @@ use miden_objects::{accounts::AccountId, Digest};
 use tracing::info;
 
 use crate::{
+    batch_builder::batch::TransactionBatch,
     block::Block,
     store::{ApplyBlock, Store},
-    SharedTxBatch, COMPONENT, MAX_NUM_CREATED_NOTES_PER_BATCH,
+    COMPONENT, MAX_NUM_CREATED_NOTES_PER_BATCH,
 };
 
 pub mod errors;
@@ -33,7 +34,7 @@ pub trait BlockBuilder: Send + Sync + 'static {
     /// block. In other words, if `build_block()` is never called, then no blocks are produced.
     async fn build_block(
         &self,
-        batches: Vec<SharedTxBatch>,
+        batches: &[TransactionBatch],
     ) -> Result<(), BuildBlockError>;
 }
 
@@ -69,7 +70,7 @@ where
 {
     async fn build_block(
         &self,
-        batches: Vec<SharedTxBatch>,
+        batches: &[TransactionBatch],
     ) -> Result<(), BuildBlockError> {
         let account_updates: Vec<(AccountId, Digest)> =
             batches.iter().flat_map(|batch| batch.updated_accounts()).collect();

block-producer/src/block_builder/prover/block_witness.rs

@@ -11,7 +11,7 @@ use miden_vm::{crypto::MerklePath, AdviceInputs, StackInputs};
 
 use crate::{
     block_builder::errors::{BlockProverError, BuildBlockError},
-    SharedTxBatch, CREATED_NOTES_SMT_DEPTH,
+    TransactionBatch, CREATED_NOTES_SMT_DEPTH,
 };
 
 // CONSTANTS
@@ -43,9 +43,9 @@ pub struct BlockWitness {
 impl BlockWitness {
     pub fn new(
         block_inputs: BlockInputs,
-        batches: Vec<SharedTxBatch>,
+        batches: &[TransactionBatch],
     ) -> Result<Self, BuildBlockError> {
-        Self::validate_inputs(&block_inputs, &batches)?;
+        Self::validate_inputs(&block_inputs, batches)?;
 
         let updated_accounts = {
             let mut account_initial_states: BTreeMap<AccountId, Digest> =
@@ -117,10 +117,11 @@ impl BlockWitness {
             // Notes stack inputs
             {
                 let num_created_notes_roots = self.batch_created_notes_roots.len();
-                for (batch_index, batch_created_notes_root) in self.batch_created_notes_roots {
-                    stack_inputs.extend(batch_created_notes_root);
+                for (batch_index, batch_created_notes_root) in self.batch_created_notes_roots.iter()
+                {
+                    stack_inputs.extend(batch_created_notes_root.iter());
 
-                    let batch_index = u64::try_from(batch_index)
+                    let batch_index = u64::try_from(*batch_index)
                         .expect("can't be more than 2^64 - 1 notes created");
                     stack_inputs.push(Felt::from(batch_index));
                 }
@@ -181,7 +182,7 @@ impl BlockWitness {
 
     fn validate_inputs(
         block_inputs: &BlockInputs,
-        batches: &[SharedTxBatch],
+        batches: &[TransactionBatch],
     ) -> Result<(), BuildBlockError> {
         // TODO:
         // - Block height returned for each nullifier is 0.
@@ -199,7 +200,7 @@ impl BlockWitness {
     /// states returned from the store
     fn validate_account_states(
         block_inputs: &BlockInputs,
-        batches: &[SharedTxBatch],
+        batches: &[TransactionBatch],
     ) -> Result<(), BuildBlockError> {
         let batches_initial_states: BTreeMap<AccountId, Digest> =
             batches.iter().flat_map(|batch| batch.account_initial_states()).collect();