feat(block-producer): multiple inflight account txs (#407)

* feat(block-producer): multiple inflight account txs

Allow multiple inflight transactions to affect the same account.

This required adding support in the transaction batch, block witness
and block producer functions to ensure that account state transitions
are correct.

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();