feat(block-producer): multiple inflight account txs

CHANGELOG.md

@@ -13,6 +13,7 @@
 - Now accounts for genesis are optional. Accounts directory will be overwritten, if `--force` flag is set (#420).
 - Added `GetAccountStateDelta` endpoint (#418).
 - Added `CheckNullifiersByPrefix` endpoint (#419).
+- Support multiple inflight transactions on the same account (#407).
 
 ### Fixes
 

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

@@ -1,19 +1,18 @@
 use std::{
-    collections::{BTreeMap, BTreeSet},
+    collections::{btree_map::Entry, BTreeMap, BTreeSet},
     mem,
 };
 
 use miden_objects::{
-    accounts::AccountId,
+    accounts::{delta::AccountUpdateDetails, AccountId},
     batches::BatchNoteTree,
-    block::BlockAccountUpdate,
     crypto::{
         hash::blake::{Blake3Digest, Blake3_256},
         merkle::MerklePath,
     },
     notes::{NoteHeader, NoteId, Nullifier},
-    transaction::{InputNoteCommitment, OutputNote, TransactionId, TxAccountUpdate},
-    Digest, MAX_NOTES_PER_BATCH,
+    transaction::{InputNoteCommitment, OutputNote, TransactionId},
+    AccountDeltaError, Digest, MAX_NOTES_PER_BATCH,
 };
 use tracing::instrument;
 
@@ -31,12 +30,45 @@ pub type BatchId = Blake3Digest<32>;
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TransactionBatch {
     id: BatchId,
-    updated_accounts: Vec<(TransactionId, TxAccountUpdate)>,
+    updated_accounts: BTreeMap<AccountId, AccountUpdate>,
     input_notes: Vec<InputNoteCommitment>,
     output_notes_smt: BatchNoteTree,
     output_notes: Vec<OutputNote>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AccountUpdate {
+    pub init_state: Digest,
+    pub final_state: Digest,
+    pub transactions: Vec<TransactionId>,
+    pub details: AccountUpdateDetails,
+}
+
+impl AccountUpdate {
+    fn new(tx: &ProvenTransaction) -> Self {
+        Self {
+            init_state: tx.account_update().init_state_hash(),
+            final_state: tx.account_update().final_state_hash(),
+            transactions: vec![tx.id()],
+            details: tx.account_update().details().clone(),
+        }
+    }
+
+    /// Merges the transaction's update into this account update.
+    fn merge_tx(&mut self, tx: &ProvenTransaction) -> Result<(), AccountDeltaError> {
+        assert!(
+            self.final_state == tx.account_update().init_state_hash(),
+            "Transacion's initial state does not match current account state"
+        );
+
+        self.final_state = tx.account_update().final_state_hash();
+        self.transactions.push(tx.id());
+        self.details = self.details.clone().merge(tx.account_update().details().clone())?;
+
+        Ok(())
+    }
+}
+
 impl TransactionBatch {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
@@ -62,14 +94,23 @@ impl TransactionBatch {
 
         // Populate batch output notes and updated accounts.
         let mut output_notes = OutputNoteTracker::new(&txs)?;
-        let mut updated_accounts = vec![];
+        let mut updated_accounts = BTreeMap::<AccountId, AccountUpdate>::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
-            //       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()));
+            // Merge account updates so that state transitions A->B->C become A->C.
+            match updated_accounts.entry(tx.account_id()) {
+                Entry::Vacant(vacant) => {
+                    vacant.insert(AccountUpdate::new(tx));
+                },
+                Entry::Occupied(occupied) => occupied.into_mut().merge_tx(tx).map_err(|error| {
+                    BuildBatchError::AccountUpdateError {
+                        account_id: tx.account_id(),
+                        error,
+                        txs: txs.clone(),
+                    }
+                })?,
+            };
+
             // Check unauthenticated input notes for duplicates:
             for note in tx.get_unauthenticated_notes() {
                 let id = note.id();
@@ -142,20 +183,13 @@ impl TransactionBatch {
     pub fn account_initial_states(&self) -> impl Iterator<Item = (AccountId, Digest)> + '_ {
         self.updated_accounts
             .iter()
-            .map(|(_, update)| (update.account_id(), update.init_state_hash()))
+            .map(|(&account_id, update)| (account_id, update.init_state))
     }
 
     /// Returns an iterator over (account_id, details, new_state_hash) tuples for accounts that were
     /// modified in this transaction batch.
-    pub fn updated_accounts(&self) -> impl Iterator<Item = BlockAccountUpdate> + '_ {
-        self.updated_accounts.iter().map(|(transaction_id, update)| {
-            BlockAccountUpdate::new(
-                update.account_id(),
-                update.final_state_hash(),
-                update.details().clone(),
-                vec![*transaction_id],
-            )
-        })
+    pub fn updated_accounts(&self) -> impl Iterator<Item = (&AccountId, &AccountUpdate)> + '_ {
+        self.updated_accounts.iter()
     }
 
     /// Returns input notes list consumed by the transactions in this batch. Any unauthenticated

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

@@ -162,12 +162,15 @@ async fn test_batch_builder_find_dangling_notes() {
         },
     ));
 
+    // An account with 5 states so that we can simulate running 2 transactions against it.
+    let account = MockPrivateAccount::<3>::from(1);
+
     let note_1 = mock_note(1);
     let note_2 = mock_note(2);
-    let tx1 = MockProvenTxBuilder::with_account_index(1)
+    let tx1 = MockProvenTxBuilder::with_account(account.id, account.states[0], account.states[1])
         .output_notes(vec![OutputNote::Full(note_1.clone())])
         .build();
-    let tx2 = MockProvenTxBuilder::with_account_index(1)
+    let tx2 = MockProvenTxBuilder::with_account(account.id, account.states[1], account.states[2])
         .unauthenticated_notes(vec![note_1.clone()])
         .output_notes(vec![OutputNote::Full(note_2.clone())])
         .build();
@@ -184,10 +187,10 @@ async fn test_batch_builder_find_dangling_notes() {
 
     let note_3 = mock_note(3);
 
-    let tx1 = MockProvenTxBuilder::with_account_index(1)
+    let tx1 = MockProvenTxBuilder::with_account(account.id, account.states[0], account.states[1])
         .unauthenticated_notes(vec![note_2.clone()])
         .build();
-    let tx2 = MockProvenTxBuilder::with_account_index(1)
+    let tx2 = MockProvenTxBuilder::with_account(account.id, account.states[1], account.states[2])
         .unauthenticated_notes(vec![note_3.clone()])
         .build();
 

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

@@ -3,7 +3,8 @@ 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},
+    accounts::AccountId,
+    block::Block,
     notes::{NoteHeader, Nullifier},
     transaction::InputNoteCommitment,
 };
@@ -74,8 +75,11 @@ where
             batches = %format_array(batches.iter().map(|batch| format_blake3_digest(batch.id()))),
         );
 
-        let updated_accounts: Vec<_> =
-            batches.iter().flat_map(TransactionBatch::updated_accounts).collect();
+        let updated_account_set: BTreeSet<AccountId> = batches
+            .iter()
+            .flat_map(TransactionBatch::updated_accounts)
+            .map(|(account_id, _)| *account_id)
+            .collect();
 
         let output_notes: Vec<_> =
             batches.iter().map(TransactionBatch::output_notes).cloned().collect();
@@ -103,7 +107,7 @@ where
         let block_inputs = self
             .store
             .get_block_inputs(
-                updated_accounts.iter().map(BlockAccountUpdate::account_id),
+                updated_account_set.into_iter(),
                 produced_nullifiers.iter(),
                 dangling_notes.iter(),
             )
@@ -117,7 +121,7 @@ where
             return Err(BuildBlockError::UnauthenticatedNotesNotFound(missing_notes));
         }
 
-        let block_header_witness = BlockWitness::new(block_inputs, batches)?;
+        let (block_header_witness, updated_accounts) = BlockWitness::new(block_inputs, batches)?;
 
         let new_block_header = self.block_kernel.prove(block_header_witness)?;
         let block_num = new_block_header.block_num();