From 91684e814dfdbe1d91bfeed5cf42f74dd0943e4e Mon Sep 17 00:00:00 2001 From: Ignacio Amigo Date: Fri, 19 Jun 2026 18:38:30 -0300 Subject: [PATCH 1/6] refactor(batch): use partaialaccount instead of account --- .../src/transaction/batch/data_store.rs | 68 +++++++++++++------ 1 file changed, 48 insertions(+), 20 deletions(-) diff --git a/crates/rust-client/src/transaction/batch/data_store.rs b/crates/rust-client/src/transaction/batch/data_store.rs index 5e44890c91..4fdd6aa725 100644 --- a/crates/rust-client/src/transaction/batch/data_store.rs +++ b/crates/rust-client/src/transaction/batch/data_store.rs @@ -6,11 +6,11 @@ use miden_protocol::account::{ Account, AccountId, PartialAccount, + PartialStorage, StorageMapKey, StorageMapWitness, - StorageSlotContent, }; -use miden_protocol::asset::{AssetVaultKey, AssetWitness}; +use miden_protocol::asset::{AssetVaultKey, AssetWitness, PartialVault}; use miden_protocol::block::{BlockHeader, BlockNumber}; use miden_protocol::note::{NoteScript, NoteScriptRoot}; use miden_protocol::transaction::{AccountInputs, PartialBlockchain}; @@ -24,32 +24,38 @@ use crate::store::data_store::ClientDataStore; // ================================================================================================ /// A [`DataStore`] that lets a [`crate::transaction::BatchBuilder`] stack in-memory account -/// states for any number of local accounts. For each account registered in -/// `current_accounts`, the executor sees the in-batch state instead of the stale store state. -/// All other reads pass through to the inner [`ClientDataStore`]. +/// inputs for any number of local accounts. For each account registered in `current_accounts`, +/// the executor sees the in-batch partial account state instead of the stale store state. All +/// other reads pass through to the inner [`ClientDataStore`]. pub(crate) struct InMemoryBatchDataStore { inner: ClientDataStore, - current_accounts: BTreeMap, + current_accounts: BTreeMap, + execution_accounts: BTreeMap, } impl InMemoryBatchDataStore { /// Wraps the provided [`ClientDataStore`] with an empty in-batch account cache. pub(crate) fn new(inner: ClientDataStore) -> Self { - Self { inner, current_accounts: BTreeMap::new() } + Self { + inner, + current_accounts: BTreeMap::new(), + execution_accounts: BTreeMap::new(), + } } - /// Returns the in-batch account state for `id`, if a transaction earlier in the batch - /// has cached one. A return of `None` means subsequent transactions targeting this - /// account will see the store's state instead. + /// Returns the full in-batch account state for `id`, if a transaction earlier in the + /// batch has cached one. A return of `None` means subsequent transactions targeting + /// this account will load the store's state first. pub(crate) fn get_account(&self, id: AccountId) -> Option<&Account> { - self.current_accounts.get(&id) + self.execution_accounts.get(&id) } /// Records the post-execution state of an account so that later transactions in the /// same batch targeting `id` observe the in-batch state. Overwrites any previously /// cached entry for `id`. pub(crate) fn cache_account(&mut self, id: AccountId, new_state: Account) { - self.current_accounts.insert(id, new_state); + self.current_accounts.insert(id, full_partial_account(&new_state)); + self.execution_accounts.insert(id, new_state); } /// Returns the inner [`ClientDataStore`]'s MAST store so callers can load account @@ -74,6 +80,18 @@ impl InMemoryBatchDataStore { } } +fn full_partial_account(account: &Account) -> PartialAccount { + PartialAccount::new( + account.id(), + account.nonce(), + account.code().clone(), + PartialStorage::new_full(account.storage().clone()), + PartialVault::new_full(account.vault().clone()), + account.seed(), + ) + .expect("account should ensure that seed is valid for account") +} + // DATA STORE IMPL // ================================================================================================ @@ -87,7 +105,7 @@ impl DataStore for InMemoryBatchDataStore { self.inner.get_transaction_inputs(account_id, ref_blocks).await?; if let Some(account) = self.current_accounts.get(&account_id) { - partial_account = PartialAccount::from(account); + partial_account = account.clone(); } Ok((partial_account, block_header, partial_blockchain)) @@ -107,8 +125,16 @@ impl DataStore for InMemoryBatchDataStore { "vault root mismatch for account {account_id}: in-batch root = {in_batch_root:?}, requested root = {vault_root:?}", ))); } - let witnesses = vault_keys.into_iter().map(|key| vault.open(key)).collect(); - Ok(witnesses) + vault_keys + .into_iter() + .map(|key| { + vault.open(key).map_err(|err| { + DataStoreError::other(format!( + "vault asset witness not found in in-batch account state for account {account_id}: {err}", + )) + }) + }) + .collect() } else { self.inner.get_vault_asset_witnesses(account_id, vault_root, vault_keys).await } @@ -121,11 +147,13 @@ impl DataStore for InMemoryBatchDataStore { map_key: StorageMapKey, ) -> Result { if let Some(account) = self.current_accounts.get(&account_id) { - for slot in account.storage().slots() { - if let StorageSlotContent::Map(map) = slot.content() - && map.root() == map_root - { - return Ok(map.open(&map_key)); + for map in account.storage().maps() { + if map.root() == map_root { + return map.open(&map_key).map_err(|err| { + DataStoreError::other(format!( + "storage map witness not found in in-batch account state for account {account_id}: {err}", + )) + }); } } return Err(DataStoreError::other(format!( From 82fc2750bf5c683f7ce24d1855c671040ff8983f Mon Sep 17 00:00:00 2001 From: Ignacio Amigo Date: Fri, 19 Jun 2026 19:21:01 -0300 Subject: [PATCH 2/6] refactor(batch): use partaialaccount instead of account --- .../rust-client/src/store/data_store/mod.rs | 15 +++ .../src/transaction/batch/data_store.rs | 113 ++++++++---------- .../rust-client/src/transaction/batch/mod.rs | 34 ++---- crates/rust-client/src/transaction/mod.rs | 21 +++- 4 files changed, 101 insertions(+), 82 deletions(-) diff --git a/crates/rust-client/src/store/data_store/mod.rs b/crates/rust-client/src/store/data_store/mod.rs index dbad8212ce..6fd2b75d57 100644 --- a/crates/rust-client/src/store/data_store/mod.rs +++ b/crates/rust-client/src/store/data_store/mod.rs @@ -63,6 +63,21 @@ impl ClientDataStore { self.cache.mast_store.clone() } + pub(crate) async fn load_account( + &self, + account_id: AccountId, + ) -> Result { + let account_record = self + .store + .get_account(account_id) + .await? + .ok_or(DataStoreError::AccountNotFound(account_id))?; + + account_record + .try_into() + .map_err(|_| DataStoreError::AccountNotFound(account_id)) + } + /// Stores the provided foreign account inputs so they can be served to the executor upon /// request. pub fn register_foreign_account_inputs( diff --git a/crates/rust-client/src/transaction/batch/data_store.rs b/crates/rust-client/src/transaction/batch/data_store.rs index 4fdd6aa725..abba1a1a62 100644 --- a/crates/rust-client/src/transaction/batch/data_store.rs +++ b/crates/rust-client/src/transaction/batch/data_store.rs @@ -4,14 +4,16 @@ use alloc::vec::Vec; use miden_protocol::account::{ Account, + AccountDelta, AccountId, PartialAccount, - PartialStorage, StorageMapKey, StorageMapWitness, + StorageSlotContent, }; -use miden_protocol::asset::{AssetVaultKey, AssetWitness, PartialVault}; +use miden_protocol::asset::{AssetVaultKey, AssetWitness}; use miden_protocol::block::{BlockHeader, BlockNumber}; +use miden_protocol::errors::AccountDeltaError; use miden_protocol::note::{NoteScript, NoteScriptRoot}; use miden_protocol::transaction::{AccountInputs, PartialBlockchain}; use miden_protocol::vm::FutureMaybeSend; @@ -24,38 +26,34 @@ use crate::store::data_store::ClientDataStore; // ================================================================================================ /// A [`DataStore`] that lets a [`crate::transaction::BatchBuilder`] stack in-memory account -/// inputs for any number of local accounts. For each account registered in `current_accounts`, -/// the executor sees the in-batch partial account state instead of the stale store state. All -/// other reads pass through to the inner [`ClientDataStore`]. +/// inputs for any number of local accounts. For each account pushed into the batch, an +/// accumulated [`AccountDelta`] is kept; the executor sees the in-batch account state (the +/// persisted account with that delta applied) instead of the stale store state. All other +/// reads pass through to the inner [`ClientDataStore`]. pub(crate) struct InMemoryBatchDataStore { inner: ClientDataStore, - current_accounts: BTreeMap, - execution_accounts: BTreeMap, + account_deltas: BTreeMap, } impl InMemoryBatchDataStore { /// Wraps the provided [`ClientDataStore`] with an empty in-batch account cache. pub(crate) fn new(inner: ClientDataStore) -> Self { - Self { - inner, - current_accounts: BTreeMap::new(), - execution_accounts: BTreeMap::new(), - } - } - - /// Returns the full in-batch account state for `id`, if a transaction earlier in the - /// batch has cached one. A return of `None` means subsequent transactions targeting - /// this account will load the store's state first. - pub(crate) fn get_account(&self, id: AccountId) -> Option<&Account> { - self.execution_accounts.get(&id) + Self { inner, account_deltas: BTreeMap::new() } } - /// Records the post-execution state of an account so that later transactions in the - /// same batch targeting `id` observe the in-batch state. Overwrites any previously - /// cached entry for `id`. - pub(crate) fn cache_account(&mut self, id: AccountId, new_state: Account) { - self.current_accounts.insert(id, full_partial_account(&new_state)); - self.execution_accounts.insert(id, new_state); + /// Merges `delta` into the accumulated in-batch delta for `account_id`, so later + /// transactions in the same batch targeting `account_id` observe its post-state. + pub(crate) fn cache_account( + &mut self, + account_id: AccountId, + delta: AccountDelta, + ) -> Result<(), AccountDeltaError> { + if let Some(existing) = self.account_deltas.get_mut(&account_id) { + existing.merge(delta)?; + } else { + self.account_deltas.insert(account_id, delta); + } + Ok(()) } /// Returns the inner [`ClientDataStore`]'s MAST store so callers can load account @@ -78,18 +76,20 @@ impl InMemoryBatchDataStore { pub(crate) fn register_note_scripts(&self, note_scripts: impl IntoIterator) { self.inner.register_note_scripts(note_scripts); } -} -fn full_partial_account(account: &Account) -> PartialAccount { - PartialAccount::new( - account.id(), - account.nonce(), - account.code().clone(), - PartialStorage::new_full(account.storage().clone()), - PartialVault::new_full(account.vault().clone()), - account.seed(), - ) - .expect("account should ensure that seed is valid for account") + pub(crate) async fn current_account( + &self, + account_id: AccountId, + ) -> Result { + let mut account = self.inner.load_account(account_id).await?; + if let Some(delta) = self.account_deltas.get(&account_id) { + account.apply_delta(delta).map_err(|err| { + DataStoreError::other_with_source("failed to apply in-batch account delta", err) + })?; + } + + Ok(account) + } } // DATA STORE IMPL @@ -104,8 +104,8 @@ impl DataStore for InMemoryBatchDataStore { let (mut partial_account, block_header, partial_blockchain) = self.inner.get_transaction_inputs(account_id, ref_blocks).await?; - if let Some(account) = self.current_accounts.get(&account_id) { - partial_account = account.clone(); + if self.account_deltas.contains_key(&account_id) { + partial_account = PartialAccount::from(&self.current_account(account_id).await?); } Ok((partial_account, block_header, partial_blockchain)) @@ -117,7 +117,8 @@ impl DataStore for InMemoryBatchDataStore { vault_root: Word, vault_keys: BTreeSet, ) -> Result, DataStoreError> { - if let Some(account) = self.current_accounts.get(&account_id) { + if self.account_deltas.contains_key(&account_id) { + let account = self.current_account(account_id).await?; let vault = account.vault(); let in_batch_root = vault.root(); if in_batch_root != vault_root { @@ -125,16 +126,8 @@ impl DataStore for InMemoryBatchDataStore { "vault root mismatch for account {account_id}: in-batch root = {in_batch_root:?}, requested root = {vault_root:?}", ))); } - vault_keys - .into_iter() - .map(|key| { - vault.open(key).map_err(|err| { - DataStoreError::other(format!( - "vault asset witness not found in in-batch account state for account {account_id}: {err}", - )) - }) - }) - .collect() + let witnesses = vault_keys.into_iter().map(|key| vault.open(key)).collect(); + Ok(witnesses) } else { self.inner.get_vault_asset_witnesses(account_id, vault_root, vault_keys).await } @@ -146,21 +139,21 @@ impl DataStore for InMemoryBatchDataStore { map_root: Word, map_key: StorageMapKey, ) -> Result { - if let Some(account) = self.current_accounts.get(&account_id) { - for map in account.storage().maps() { - if map.root() == map_root { - return map.open(&map_key).map_err(|err| { - DataStoreError::other(format!( - "storage map witness not found in in-batch account state for account {account_id}: {err}", - )) - }); + if self.account_deltas.contains_key(&account_id) { + let account = self.current_account(account_id).await?; + for slot in account.storage().slots() { + if let StorageSlotContent::Map(map) = slot.content() + && map.root() == map_root + { + return Ok(map.open(&map_key)); } } - return Err(DataStoreError::other(format!( + Err(DataStoreError::other(format!( "storage map root not found in in-batch account state for account {account_id}: requested root = {map_root:?}", - ))); + ))) + } else { + self.inner.get_storage_map_witness(account_id, map_root, map_key).await } - self.inner.get_storage_map_witness(account_id, map_root, map_key).await } async fn get_foreign_account_inputs( diff --git a/crates/rust-client/src/transaction/batch/mod.rs b/crates/rust-client/src/transaction/batch/mod.rs index 57cd8edfa2..073f1d7c11 100644 --- a/crates/rust-client/src/transaction/batch/mod.rs +++ b/crates/rust-client/src/transaction/batch/mod.rs @@ -56,7 +56,7 @@ use alloc::vec::Vec; pub(crate) use data_store::InMemoryBatchDataStore; pub use error::BatchBuilderError; use miden_protocol::MIN_PROOF_SECURITY_LEVEL; -use miden_protocol::account::{Account, AccountId}; +use miden_protocol::account::AccountId; use miden_protocol::batch::ProposedBatch; use miden_protocol::block::{BlockHeader, BlockNumber}; use miden_protocol::note::NoteId; @@ -272,23 +272,14 @@ async fn execute_transaction_for_batch( where AUTH: TransactionAuthenticator + Sync + 'static, { - let mut account = if let Some(account) = data_store.get_account(account_id) { - account.clone() - } else { - let record = client - .store - .get_account(account_id) - .await? - .ok_or(ClientError::AccountDataNotFound(account_id))?; - if record.is_locked() { - return Err(ClientError::AccountLocked(account_id)); - } - let account: Account = record.try_into()?; - account - }; + if client.account_reader(account_id).status().await?.is_locked() { + return Err(ClientError::AccountLocked(account_id)); + } + + let account = data_store.current_account(account_id).await?; let account_id = account.id(); - let prep = client.prepare_transaction(&account, transaction_request).await?; + let prep = client.prepare_transaction_for_batch(&account, transaction_request).await?; data_store.register_note_scripts(prep.output_note_scripts()); for fpi_account in &prep.foreign_account_inputs { @@ -309,12 +300,13 @@ where .build_executor(data_store)? .execute_transaction(account_id, prep.block_num, notes, prep.tx_args) .await?; + let account_delta = executed_transaction.account_delta().clone(); - // Cache new account state in memory data store - account - .apply_delta(executed_transaction.account_delta()) - .map_err(ClientError::AccountError)?; - data_store.cache_account(account_id, account); + // Merge this transaction's delta into the batch's in-memory state so later pushes for this + // account observe its post-state. + data_store + .cache_account(account_id, account_delta) + .map_err(|err| ClientError::StoreError(crate::store::StoreError::AccountDeltaError(err)))?; validate_executed_transaction(&executed_transaction, &prep.output_recipients)?; TransactionResult::new(executed_transaction, prep.future_notes) diff --git a/crates/rust-client/src/transaction/mod.rs b/crates/rust-client/src/transaction/mod.rs index 7292435305..d542fbc90d 100644 --- a/crates/rust-client/src/transaction/mod.rs +++ b/crates/rust-client/src/transaction/mod.rs @@ -348,10 +348,29 @@ where &self, account: &Account, transaction_request: TransactionRequest, + ) -> Result { + self.prepare_transaction_inner(account, transaction_request, true).await + } + + pub(crate) async fn prepare_transaction_for_batch( + &self, + account: &Account, + transaction_request: TransactionRequest, + ) -> Result { + self.prepare_transaction_inner(account, transaction_request, false).await + } + + async fn prepare_transaction_inner( + &self, + account: &Account, + transaction_request: TransactionRequest, + validate_account: bool, ) -> Result { let account_id = account.id(); self.validate_recency().await?; - validate_account_request(&transaction_request, account)?; + if validate_account { + validate_account_request(&transaction_request, account)?; + } // Retrieve all input notes from the store. let mut stored_note_records = self From 8837ea81e26916f3fb2f8350ac0f72345ad27d45 Mon Sep 17 00:00:00 2001 From: Ignacio Amigo Date: Tue, 23 Jun 2026 17:04:15 -0300 Subject: [PATCH 3/6] refactor: remove from basic check: --- .../rust-client/src/store/data_store/mod.rs | 6 +- .../src/transaction/batch/data_store.rs | 102 +++++----- .../rust-client/src/transaction/batch/mod.rs | 180 ++++++++++++++++-- crates/rust-client/src/transaction/mod.rs | 52 +++-- 4 files changed, 244 insertions(+), 96 deletions(-) diff --git a/crates/rust-client/src/store/data_store/mod.rs b/crates/rust-client/src/store/data_store/mod.rs index 6fd2b75d57..348d32320f 100644 --- a/crates/rust-client/src/store/data_store/mod.rs +++ b/crates/rust-client/src/store/data_store/mod.rs @@ -63,13 +63,13 @@ impl ClientDataStore { self.cache.mast_store.clone() } - pub(crate) async fn load_account( + pub(crate) async fn load_partial_account( &self, account_id: AccountId, - ) -> Result { + ) -> Result { let account_record = self .store - .get_account(account_id) + .get_minimal_partial_account(account_id) .await? .ok_or(DataStoreError::AccountNotFound(account_id))?; diff --git a/crates/rust-client/src/transaction/batch/data_store.rs b/crates/rust-client/src/transaction/batch/data_store.rs index abba1a1a62..4854be908f 100644 --- a/crates/rust-client/src/transaction/batch/data_store.rs +++ b/crates/rust-client/src/transaction/batch/data_store.rs @@ -2,20 +2,11 @@ use alloc::collections::{BTreeMap, BTreeSet}; use alloc::sync::Arc; use alloc::vec::Vec; -use miden_protocol::account::{ - Account, - AccountDelta, - AccountId, - PartialAccount, - StorageMapKey, - StorageMapWitness, - StorageSlotContent, -}; +use miden_protocol::account::{AccountId, PartialAccount, StorageMapKey, StorageMapWitness}; use miden_protocol::asset::{AssetVaultKey, AssetWitness}; use miden_protocol::block::{BlockHeader, BlockNumber}; -use miden_protocol::errors::AccountDeltaError; use miden_protocol::note::{NoteScript, NoteScriptRoot}; -use miden_protocol::transaction::{AccountInputs, PartialBlockchain}; +use miden_protocol::transaction::{AccountInputs, PartialBlockchain, TransactionInputs}; use miden_protocol::vm::FutureMaybeSend; use miden_protocol::{MastForest, Word}; use miden_tx::{DataStore, DataStoreError, MastForestStore, TransactionMastStore}; @@ -26,34 +17,31 @@ use crate::store::data_store::ClientDataStore; // ================================================================================================ /// A [`DataStore`] that lets a [`crate::transaction::BatchBuilder`] stack in-memory account -/// inputs for any number of local accounts. For each account pushed into the batch, an -/// accumulated [`AccountDelta`] is kept; the executor sees the in-batch account state (the -/// persisted account with that delta applied) instead of the stale store state. All other -/// reads pass through to the inner [`ClientDataStore`]. +/// inputs for any number of local accounts. For each account pushed into the batch, a +/// [`PartialAccount`] is cached; the executor sees the in-batch partial account state instead of +/// the stale store state. Witness reads for the cached account are resolved from the prior +/// transaction's execution advice, without reconstructing a full account. pub(crate) struct InMemoryBatchDataStore { inner: ClientDataStore, - account_deltas: BTreeMap, + current_accounts: BTreeMap, +} + +struct CachedAccountState { + account: PartialAccount, + tx_inputs: TransactionInputs, } impl InMemoryBatchDataStore { /// Wraps the provided [`ClientDataStore`] with an empty in-batch account cache. pub(crate) fn new(inner: ClientDataStore) -> Self { - Self { inner, account_deltas: BTreeMap::new() } + Self { inner, current_accounts: BTreeMap::new() } } - /// Merges `delta` into the accumulated in-batch delta for `account_id`, so later - /// transactions in the same batch targeting `account_id` observe its post-state. - pub(crate) fn cache_account( - &mut self, - account_id: AccountId, - delta: AccountDelta, - ) -> Result<(), AccountDeltaError> { - if let Some(existing) = self.account_deltas.get_mut(&account_id) { - existing.merge(delta)?; - } else { - self.account_deltas.insert(account_id, delta); - } - Ok(()) + /// Caches the post-transaction partial account and the transaction inputs carrying the + /// execution advice needed to serve later witness reads for this in-batch state. + pub(crate) fn cache_account(&mut self, account: PartialAccount, tx_inputs: TransactionInputs) { + self.current_accounts + .insert(account.id(), CachedAccountState { account, tx_inputs }); } /// Returns the inner [`ClientDataStore`]'s MAST store so callers can load account @@ -80,15 +68,12 @@ impl InMemoryBatchDataStore { pub(crate) async fn current_account( &self, account_id: AccountId, - ) -> Result { - let mut account = self.inner.load_account(account_id).await?; - if let Some(delta) = self.account_deltas.get(&account_id) { - account.apply_delta(delta).map_err(|err| { - DataStoreError::other_with_source("failed to apply in-batch account delta", err) - })?; + ) -> Result { + if let Some(state) = self.current_accounts.get(&account_id) { + return Ok(state.account.clone()); } - Ok(account) + self.inner.load_partial_account(account_id).await } } @@ -104,8 +89,8 @@ impl DataStore for InMemoryBatchDataStore { let (mut partial_account, block_header, partial_blockchain) = self.inner.get_transaction_inputs(account_id, ref_blocks).await?; - if self.account_deltas.contains_key(&account_id) { - partial_account = PartialAccount::from(&self.current_account(account_id).await?); + if let Some(state) = self.current_accounts.get(&account_id) { + partial_account = state.account.clone(); } Ok((partial_account, block_header, partial_blockchain)) @@ -117,17 +102,22 @@ impl DataStore for InMemoryBatchDataStore { vault_root: Word, vault_keys: BTreeSet, ) -> Result, DataStoreError> { - if self.account_deltas.contains_key(&account_id) { - let account = self.current_account(account_id).await?; - let vault = account.vault(); + if let Some(state) = self.current_accounts.get(&account_id) { + let vault = state.account.vault(); let in_batch_root = vault.root(); if in_batch_root != vault_root { return Err(DataStoreError::other(format!( "vault root mismatch for account {account_id}: in-batch root = {in_batch_root:?}, requested root = {vault_root:?}", ))); } - let witnesses = vault_keys.into_iter().map(|key| vault.open(key)).collect(); - Ok(witnesses) + state.tx_inputs.read_vault_asset_witnesses(vault_root, vault_keys).map_err(|err| { + DataStoreError::other_with_source( + format!( + "vault asset witness not found in in-batch account state for account {account_id}" + ), + err, + ) + }) } else { self.inner.get_vault_asset_witnesses(account_id, vault_root, vault_keys).await } @@ -139,18 +129,20 @@ impl DataStore for InMemoryBatchDataStore { map_root: Word, map_key: StorageMapKey, ) -> Result { - if self.account_deltas.contains_key(&account_id) { - let account = self.current_account(account_id).await?; - for slot in account.storage().slots() { - if let StorageSlotContent::Map(map) = slot.content() - && map.root() == map_root - { - return Ok(map.open(&map_key)); - } + if let Some(state) = self.current_accounts.get(&account_id) { + if !state.account.storage().header().map_slot_roots().any(|root| root == map_root) { + return Err(DataStoreError::other(format!( + "storage map root not found in in-batch account state for account {account_id}: requested root = {map_root:?}", + ))); } - Err(DataStoreError::other(format!( - "storage map root not found in in-batch account state for account {account_id}: requested root = {map_root:?}", - ))) + state.tx_inputs.read_storage_map_witness(map_root, map_key).map_err(|err| { + DataStoreError::other_with_source( + format!( + "storage map witness not found in in-batch account state for account {account_id}" + ), + err, + ) + }) } else { self.inner.get_storage_map_witness(account_id, map_root, map_key).await } diff --git a/crates/rust-client/src/transaction/batch/mod.rs b/crates/rust-client/src/transaction/batch/mod.rs index 073f1d7c11..138b914287 100644 --- a/crates/rust-client/src/transaction/batch/mod.rs +++ b/crates/rust-client/src/transaction/batch/mod.rs @@ -55,21 +55,25 @@ use alloc::vec::Vec; pub(crate) use data_store::InMemoryBatchDataStore; pub use error::BatchBuilderError; -use miden_protocol::MIN_PROOF_SECURITY_LEVEL; -use miden_protocol::account::AccountId; +use miden_protocol::account::{ + AccountId, AccountStorageHeader, PartialAccount, PartialStorage, PartialStorageMap, + StorageMapDelta, StorageSlotDelta, StorageSlotHeader, StorageSlotType, +}; +use miden_protocol::asset::PartialVault; use miden_protocol::batch::ProposedBatch; use miden_protocol::block::{BlockHeader, BlockNumber}; +use miden_protocol::crypto::merkle::smt::PartialSmt; use miden_protocol::note::NoteId; -use miden_protocol::transaction::{PartialBlockchain, ProvenTransaction, TransactionInputs}; -use miden_tx::auth::TransactionAuthenticator; +use miden_protocol::transaction::{ + ExecutedTransaction, PartialBlockchain, ProvenTransaction, TransactionInputs, +}; +use miden_protocol::{MIN_PROOF_SECURITY_LEVEL, ZERO}; +use miden_tx::{DataStoreError, auth::TransactionAuthenticator}; use miden_tx_batch_prover::LocalBatchProver; use crate::store::data_store::build_partial_mmr_with_paths; use crate::transaction::{ - TransactionRequest, - TransactionResult, - TransactionStoreUpdate, - validate_executed_transaction, + TransactionRequest, TransactionResult, TransactionStoreUpdate, validate_executed_transaction, }; use crate::{Client, ClientError}; @@ -261,8 +265,8 @@ where } } -/// Executes a single transaction, that is part of the batch to be sent to the node. -/// Transaction is ran as the provided `Account` +/// Executes a single transaction that is part of the batch to be sent to the node. +/// The transaction runs against the current in-batch [`PartialAccount`] state. async fn execute_transaction_for_batch( client: &Client, data_store: &mut InMemoryBatchDataStore, @@ -292,7 +296,12 @@ where let mut notes = prep.notes; if prep.ignore_invalid_notes { notes = client - .get_valid_input_notes(&account, notes, prep.tx_args.clone(), &prep.output_recipients) + .get_valid_input_notes_with_data_store( + data_store, + account_id, + notes, + prep.tx_args.clone(), + ) .await?; } @@ -300,14 +309,151 @@ where .build_executor(data_store)? .execute_transaction(account_id, prep.block_num, notes, prep.tx_args) .await?; - let account_delta = executed_transaction.account_delta().clone(); - // Merge this transaction's delta into the batch's in-memory state so later pushes for this - // account observe its post-state. - data_store - .cache_account(account_id, account_delta) - .map_err(|err| ClientError::StoreError(crate::store::StoreError::AccountDeltaError(err)))?; + let current_account = partial_account_from_executed_transaction(&executed_transaction)?; + let tx_inputs = executed_transaction.tx_inputs().clone(); + data_store.cache_account(current_account, tx_inputs); validate_executed_transaction(&executed_transaction, &prep.output_recipients)?; TransactionResult::new(executed_transaction, prep.future_notes) } + +fn partial_account_from_executed_transaction( + executed_transaction: &ExecutedTransaction, +) -> Result { + let initial_account = executed_transaction.initial_account(); + let final_account = executed_transaction.final_account(); + let code = executed_transaction + .account_delta() + .code() + .unwrap_or_else(|| initial_account.code()) + .clone(); + + if final_account.code_commitment() != code.commitment() { + return Err(DataStoreError::other(format!( + "account code commitment changed for account {} while preparing in-batch state", + final_account.id() + ))); + } + + let storage_header = final_storage_header_from_delta(executed_transaction)?; + + let storage = PartialStorage::new(storage_header, core::iter::empty::()) + .map_err(|err| { + DataStoreError::other_with_source( + "failed to rebuild final in-batch partial account storage", + err, + ) + })?; + let vault = PartialVault::new(final_account.vault_root()); + let seed = if final_account.nonce() == ZERO { + initial_account.seed() + } else { + None + }; + + PartialAccount::new(final_account.id(), final_account.nonce(), code, storage, vault, seed) + .map_err(|err| { + DataStoreError::other_with_source( + "failed to rebuild final in-batch partial account", + err, + ) + }) +} + +fn final_storage_header_from_delta( + executed_transaction: &ExecutedTransaction, +) -> Result { + let initial_account = executed_transaction.initial_account(); + let final_account = executed_transaction.final_account(); + let storage_delta = executed_transaction.account_delta().storage(); + + let mut slots = Vec::new(); + for slot in initial_account.storage().header().slots() { + let new_slot_value = match storage_delta.get(slot.name()) { + None => slot.value(), + Some(StorageSlotDelta::Value(value)) => { + if slot.slot_type() != StorageSlotType::Value { + return Err(DataStoreError::other(format!( + "storage slot {} changed as value but initial in-batch state has type {:?}", + slot.name(), + slot.slot_type() + ))); + } + *value + }, + Some(StorageSlotDelta::Map(map_delta)) => { + if slot.slot_type() != StorageSlotType::Map { + return Err(DataStoreError::other(format!( + "storage slot {} changed as map but initial in-batch state has type {:?}", + slot.name(), + slot.slot_type() + ))); + } + updated_storage_map_root(executed_transaction.tx_inputs(), slot.value(), map_delta)? + }, + }; + + slots.push(StorageSlotHeader::new( + slot.name().clone(), + slot.slot_type(), + new_slot_value, + )); + } + + let storage_header = AccountStorageHeader::new(slots).map_err(|err| { + DataStoreError::other_with_source( + "failed to rebuild final in-batch account storage header", + err, + ) + })?; + + if storage_header.to_commitment() != final_account.storage_commitment() { + return Err(DataStoreError::other(format!( + "rebuilt storage commitment does not match final account state for account {}: rebuilt = {:?}, final = {:?}", + final_account.id(), + storage_header.to_commitment(), + final_account.storage_commitment() + ))); + } + + Ok(storage_header) +} + +fn updated_storage_map_root( + tx_inputs: &TransactionInputs, + initial_root: miden_protocol::Word, + map_delta: &StorageMapDelta, +) -> Result { + let mut partial_smt = PartialSmt::new(initial_root); + + for map_key in map_delta.entries().keys() { + let witness = tx_inputs + .read_storage_map_witness(initial_root, *map_key) + .map_err(|err| { + DataStoreError::other_with_source( + "failed to read initial storage map witness while rebuilding in-batch state", + err, + ) + })?; + partial_smt.add_proof(witness.into()).map_err(|err| { + DataStoreError::other_with_source( + "failed to add storage map witness while rebuilding in-batch state", + err, + ) + })?; + } + + for (map_key, value) in map_delta.entries() { + partial_smt + .insert(map_key.hash().as_word(), *value) + .map_err(|err| { + DataStoreError::other_with_source( + "failed to apply storage map delta while rebuilding in-batch state", + err, + ) + })?; + } + + Ok(partial_smt.root()) +} diff --git a/crates/rust-client/src/transaction/mod.rs b/crates/rust-client/src/transaction/mod.rs index d542fbc90d..a3d8ddf309 100644 --- a/crates/rust-client/src/transaction/mod.rs +++ b/crates/rust-client/src/transaction/mod.rs @@ -67,7 +67,7 @@ use alloc::collections::{BTreeMap, BTreeSet}; use alloc::sync::Arc; use alloc::vec::Vec; -use miden_protocol::account::{Account, AccountCode, AccountId}; +use miden_protocol::account::{Account, AccountCode, AccountId, PartialAccount}; use miden_protocol::asset::{Asset, NonFungibleAsset}; use miden_protocol::block::BlockNumber; use miden_protocol::errors::AssetError; @@ -334,41 +334,40 @@ where } /// Performs the data-store-independent setup shared by `execute_transaction` and - /// `execute_transaction_for_batch`: validates the request against the supplied - /// `account`, loads/filters input notes, builds the transaction script and args, - /// retrieves foreign-account inputs, and computes the reference block number. + /// `execute_transaction_for_batch`: validates the request when a full account is available, + /// loads/filters input notes, builds the transaction script and args, retrieves + /// foreign-account inputs, and computes the reference block number. /// /// This method does not write to the store: any state produced by the transaction is /// persisted only after the transaction executes successfully. /// - /// `account` is the state validation runs against — for a single transaction this is - /// the persisted account; inside [`crate::transaction::BatchBuilder::push`] it is the - /// in-batch (stacked) state, so balances reflect prior pushes. + /// In batch execution, request validation that needs a full account is skipped here; the + /// executor still runs against the in-batch [`PartialAccount`] state. pub(crate) async fn prepare_transaction( &self, account: &Account, transaction_request: TransactionRequest, ) -> Result { - self.prepare_transaction_inner(account, transaction_request, true).await + self.prepare_transaction_inner(account.id(), transaction_request, Some(account)) + .await } pub(crate) async fn prepare_transaction_for_batch( &self, - account: &Account, + account: &PartialAccount, transaction_request: TransactionRequest, ) -> Result { - self.prepare_transaction_inner(account, transaction_request, false).await + self.prepare_transaction_inner(account.id(), transaction_request, None).await } async fn prepare_transaction_inner( &self, - account: &Account, + account_id: AccountId, transaction_request: TransactionRequest, - validate_account: bool, + account_to_validate: Option<&Account>, ) -> Result { - let account_id = account.id(); self.validate_recency().await?; - if validate_account { + if let Some(account) = account_to_validate { validate_account_request(&transaction_request, account)?; } @@ -681,18 +680,29 @@ where pub(crate) async fn get_valid_input_notes( &self, account: &Account, - mut input_notes: InputNotes, + input_notes: InputNotes, tx_args: TransactionArgs, output_recipients: &[NoteRecipient], ) -> Result, ClientError> { - loop { - let data_store = ClientDataStore::new(self.store.clone(), self.rpc_api.clone()); - data_store.register_note_scripts(output_recipients.iter().map(|r| r.script().clone())); + let data_store = ClientDataStore::new(self.store.clone(), self.rpc_api.clone()); + data_store.register_note_scripts(output_recipients.iter().map(|r| r.script().clone())); + data_store.mast_store().load_account_code(account.code()); - data_store.mast_store().load_account_code(account.code()); - let execution = NoteConsumptionChecker::new(&self.build_executor(&data_store)?) + self.get_valid_input_notes_with_data_store(&data_store, account.id(), input_notes, tx_args) + .await + } + + pub(crate) async fn get_valid_input_notes_with_data_store( + &self, + data_store: &STORE, + account_id: AccountId, + mut input_notes: InputNotes, + tx_args: TransactionArgs, + ) -> Result, ClientError> { + loop { + let execution = NoteConsumptionChecker::new(&self.build_executor(data_store)?) .check_notes_consumability( - account.id(), + account_id, self.store.get_sync_height().await?, input_notes.iter().map(|n| n.clone().into_note()).collect(), tx_args.clone(), From 1c297c88b9f9b0cfe9da6c638752484995bcc602 Mon Sep 17 00:00:00 2001 From: Ignacio Amigo Date: Tue, 23 Jun 2026 17:18:59 -0300 Subject: [PATCH 4/6] refactor: accountreader --- .../rust-client/src/account/account_reader.rs | 10 ++++ .../rust-client/src/store/data_store/mod.rs | 15 ------ .../src/transaction/batch/data_store.rs | 9 ++-- .../rust-client/src/transaction/batch/mod.rs | 54 +++++++++++-------- 4 files changed, 47 insertions(+), 41 deletions(-) diff --git a/crates/rust-client/src/account/account_reader.rs b/crates/rust-client/src/account/account_reader.rs index 54947fa6a4..6a8cb176c8 100644 --- a/crates/rust-client/src/account/account_reader.rs +++ b/crates/rust-client/src/account/account_reader.rs @@ -6,6 +6,7 @@ use alloc::vec::Vec; use miden_protocol::account::{ AccountHeader, AccountId, + PartialAccount, StorageMapKey, StorageMapWitness, StorageSlotName, @@ -100,6 +101,15 @@ impl AccountReader { .ok_or(ClientError::AccountDataNotFound(self.account_id)) } + /// Retrieves the minimal partial account representation for this account. + pub(crate) async fn partial_account(&self) -> Result { + self.store + .get_minimal_partial_account(self.account_id) + .await? + .ok_or(ClientError::AccountDataNotFound(self.account_id))? + .try_into() + } + /// Retrieves the addresses associated with this account. pub async fn addresses(&self) -> Result, ClientError> { self.store diff --git a/crates/rust-client/src/store/data_store/mod.rs b/crates/rust-client/src/store/data_store/mod.rs index 348d32320f..dbad8212ce 100644 --- a/crates/rust-client/src/store/data_store/mod.rs +++ b/crates/rust-client/src/store/data_store/mod.rs @@ -63,21 +63,6 @@ impl ClientDataStore { self.cache.mast_store.clone() } - pub(crate) async fn load_partial_account( - &self, - account_id: AccountId, - ) -> Result { - let account_record = self - .store - .get_minimal_partial_account(account_id) - .await? - .ok_or(DataStoreError::AccountNotFound(account_id))?; - - account_record - .try_into() - .map_err(|_| DataStoreError::AccountNotFound(account_id)) - } - /// Stores the provided foreign account inputs so they can be served to the executor upon /// request. pub fn register_foreign_account_inputs( diff --git a/crates/rust-client/src/transaction/batch/data_store.rs b/crates/rust-client/src/transaction/batch/data_store.rs index 4854be908f..279ffee5ea 100644 --- a/crates/rust-client/src/transaction/batch/data_store.rs +++ b/crates/rust-client/src/transaction/batch/data_store.rs @@ -11,6 +11,8 @@ use miden_protocol::vm::FutureMaybeSend; use miden_protocol::{MastForest, Word}; use miden_tx::{DataStore, DataStoreError, MastForestStore, TransactionMastStore}; +use crate::ClientError; +use crate::account::AccountReader; use crate::store::data_store::ClientDataStore; // IN-MEMORY BATCH DATA STORE @@ -67,13 +69,14 @@ impl InMemoryBatchDataStore { pub(crate) async fn current_account( &self, - account_id: AccountId, - ) -> Result { + account_reader: &AccountReader, + ) -> Result { + let account_id = account_reader.account_id(); if let Some(state) = self.current_accounts.get(&account_id) { return Ok(state.account.clone()); } - self.inner.load_partial_account(account_id).await + account_reader.partial_account().await } } diff --git a/crates/rust-client/src/transaction/batch/mod.rs b/crates/rust-client/src/transaction/batch/mod.rs index 138b914287..7a6732a085 100644 --- a/crates/rust-client/src/transaction/batch/mod.rs +++ b/crates/rust-client/src/transaction/batch/mod.rs @@ -56,8 +56,15 @@ use alloc::vec::Vec; pub(crate) use data_store::InMemoryBatchDataStore; pub use error::BatchBuilderError; use miden_protocol::account::{ - AccountId, AccountStorageHeader, PartialAccount, PartialStorage, PartialStorageMap, - StorageMapDelta, StorageSlotDelta, StorageSlotHeader, StorageSlotType, + AccountId, + AccountStorageHeader, + PartialAccount, + PartialStorage, + PartialStorageMap, + StorageMapDelta, + StorageSlotDelta, + StorageSlotHeader, + StorageSlotType, }; use miden_protocol::asset::PartialVault; use miden_protocol::batch::ProposedBatch; @@ -65,15 +72,22 @@ use miden_protocol::block::{BlockHeader, BlockNumber}; use miden_protocol::crypto::merkle::smt::PartialSmt; use miden_protocol::note::NoteId; use miden_protocol::transaction::{ - ExecutedTransaction, PartialBlockchain, ProvenTransaction, TransactionInputs, + ExecutedTransaction, + PartialBlockchain, + ProvenTransaction, + TransactionInputs, }; use miden_protocol::{MIN_PROOF_SECURITY_LEVEL, ZERO}; -use miden_tx::{DataStoreError, auth::TransactionAuthenticator}; +use miden_tx::DataStoreError; +use miden_tx::auth::TransactionAuthenticator; use miden_tx_batch_prover::LocalBatchProver; use crate::store::data_store::build_partial_mmr_with_paths; use crate::transaction::{ - TransactionRequest, TransactionResult, TransactionStoreUpdate, validate_executed_transaction, + TransactionRequest, + TransactionResult, + TransactionStoreUpdate, + validate_executed_transaction, }; use crate::{Client, ClientError}; @@ -276,11 +290,12 @@ async fn execute_transaction_for_batch( where AUTH: TransactionAuthenticator + Sync + 'static, { - if client.account_reader(account_id).status().await?.is_locked() { + let account_reader = client.account_reader(account_id); + if account_reader.status().await?.is_locked() { return Err(ClientError::AccountLocked(account_id)); } - let account = data_store.current_account(account_id).await?; + let account = data_store.current_account(&account_reader).await?; let account_id = account.id(); let prep = client.prepare_transaction_for_batch(&account, transaction_request).await?; @@ -394,11 +409,7 @@ fn final_storage_header_from_delta( }, }; - slots.push(StorageSlotHeader::new( - slot.name().clone(), - slot.slot_type(), - new_slot_value, - )); + slots.push(StorageSlotHeader::new(slot.name().clone(), slot.slot_type(), new_slot_value)); } let storage_header = AccountStorageHeader::new(slots).map_err(|err| { @@ -428,9 +439,8 @@ fn updated_storage_map_root( let mut partial_smt = PartialSmt::new(initial_root); for map_key in map_delta.entries().keys() { - let witness = tx_inputs - .read_storage_map_witness(initial_root, *map_key) - .map_err(|err| { + let witness = + tx_inputs.read_storage_map_witness(initial_root, *map_key).map_err(|err| { DataStoreError::other_with_source( "failed to read initial storage map witness while rebuilding in-batch state", err, @@ -445,14 +455,12 @@ fn updated_storage_map_root( } for (map_key, value) in map_delta.entries() { - partial_smt - .insert(map_key.hash().as_word(), *value) - .map_err(|err| { - DataStoreError::other_with_source( - "failed to apply storage map delta while rebuilding in-batch state", - err, - ) - })?; + partial_smt.insert(map_key.hash().as_word(), *value).map_err(|err| { + DataStoreError::other_with_source( + "failed to apply storage map delta while rebuilding in-batch state", + err, + ) + })?; } Ok(partial_smt.root()) From 7cc9017cf79b8f4f3bb592ed6d4cce0b07b00706 Mon Sep 17 00:00:00 2001 From: Ignacio Amigo Date: Thu, 25 Jun 2026 20:14:20 -0300 Subject: [PATCH 5/6] refactor: allow new keys to be read and updated --- crates/rust-client/src/errors.rs | 10 +- .../rust-client/src/store/data_store/mod.rs | 6 + crates/rust-client/src/store/errors.rs | 2 + crates/rust-client/src/store/mod.rs | 43 ++++ crates/rust-client/src/store/smt_forest.rs | 58 +++++ .../src/transaction/batch/data_store.rs | 141 +++++++--- .../rust-client/src/transaction/batch/mod.rs | 3 +- crates/sqlite-store/src/account/accounts.rs | 47 ++++ crates/sqlite-store/src/account/vault.rs | 109 ++++++-- crates/sqlite-store/src/lib.rs | 43 ++++ .../miden-client-tests/src/tests/batch.rs | 242 +++++++++++++++++- 11 files changed, 634 insertions(+), 70 deletions(-) diff --git a/crates/rust-client/src/errors.rs b/crates/rust-client/src/errors.rs index 41f7504558..4ff6259775 100644 --- a/crates/rust-client/src/errors.rs +++ b/crates/rust-client/src/errors.rs @@ -6,7 +6,13 @@ use core::fmt; use miden_protocol::Word; use miden_protocol::account::AccountId; use miden_protocol::crypto::merkle::MerkleError; -pub use miden_protocol::errors::{AccountError, AccountIdError, AssetError, NetworkIdError}; +pub use miden_protocol::errors::{ + AccountDeltaError, + AccountError, + AccountIdError, + AssetError, + NetworkIdError, +}; use miden_protocol::errors::{ NoteError, PartialBlockchainError, @@ -78,6 +84,8 @@ pub enum ClientError { AccountAlreadyTracked(AccountId), #[error("account error")] AccountError(#[from] AccountError), + #[error("account delta error")] + AccountDeltaError(#[from] AccountDeltaError), #[error("account {0} is locked because the local state may be out of date with the network")] AccountLocked(AccountId), #[error( diff --git a/crates/rust-client/src/store/data_store/mod.rs b/crates/rust-client/src/store/data_store/mod.rs index dbad8212ce..3a04e45055 100644 --- a/crates/rust-client/src/store/data_store/mod.rs +++ b/crates/rust-client/src/store/data_store/mod.rs @@ -63,6 +63,12 @@ impl ClientDataStore { self.cache.mast_store.clone() } + /// Returns the underlying [`Store`], so callers (e.g. the batch builder) can serve witnesses + /// against in-batch state without reconstructing accounts. + pub fn store(&self) -> &Arc { + &self.store + } + /// Stores the provided foreign account inputs so they can be served to the executor upon /// request. pub fn register_foreign_account_inputs( diff --git a/crates/rust-client/src/store/errors.rs b/crates/rust-client/src/store/errors.rs index c4b7e00217..4a71793dda 100644 --- a/crates/rust-client/src/store/errors.rs +++ b/crates/rust-client/src/store/errors.rs @@ -92,6 +92,8 @@ pub enum StoreError { VaultDataNotFound(Word), #[error("failed to parse word")] WordError(#[from] WordError), + #[error("operation `{0}` is not supported by this store backend")] + UnsupportedOperation(&'static str), } impl From for DataStoreError { diff --git a/crates/rust-client/src/store/mod.rs b/crates/rust-client/src/store/mod.rs index 3a2ed543ec..48a6f152d9 100644 --- a/crates/rust-client/src/store/mod.rs +++ b/crates/rust-client/src/store/mod.rs @@ -29,6 +29,7 @@ use core::fmt::Debug; use miden_protocol::account::{ Account, AccountCode, + AccountDelta, AccountHeader, AccountId, AccountStorage, @@ -674,6 +675,48 @@ pub trait Store: Send + Sync { } } + // IN-BATCH (STAGED) WITNESSES + // -------------------------------------------------------------------------------------------- + + /// Returns vault asset witnesses for `vault_keys` as the account's vault would look after + /// applying `delta` to its committed state, *without* persisting the change. `vault_root` is + /// the in-batch vault root the witnesses are expected to be valid against. + /// + /// This lets [`crate::transaction::BatchBuilder`] serve witnesses for keys a prior in-batch + /// transaction never touched (and so are absent from its execution advice) against the + /// stacked account state, without reconstructing the full account. + /// + /// The default implementation returns [`StoreError::UnsupportedOperation`]; backends that + /// keep an in-memory Merkle forest (e.g. `SqliteStore`) override it by staging the delta on + /// that forest. + async fn vault_asset_witnesses_after_delta( + &self, + account_id: AccountId, + delta: AccountDelta, + vault_root: Word, + vault_keys: BTreeSet, + ) -> Result, StoreError> { + let _ = (account_id, delta, vault_root, vault_keys); + Err(StoreError::UnsupportedOperation("vault_asset_witnesses_after_delta")) + } + + /// Returns the storage map witness for `map_key` as the account's storage would look after + /// applying `delta` to its committed state, *without* persisting the change. `map_root` is + /// the in-batch root of the map slot the witness is expected to be valid against. + /// + /// See [`Store::vault_asset_witnesses_after_delta`] for the rationale and the default + /// behavior. + async fn storage_map_witness_after_delta( + &self, + account_id: AccountId, + delta: AccountDelta, + map_root: Word, + map_key: StorageMapKey, + ) -> Result { + let _ = (account_id, delta, map_root, map_key); + Err(StoreError::UnsupportedOperation("storage_map_witness_after_delta")) + } + // PARTIAL ACCOUNTS // -------------------------------------------------------------------------------------------- diff --git a/crates/rust-client/src/store/smt_forest.rs b/crates/rust-client/src/store/smt_forest.rs index c70af10f85..ce52f8d0bf 100644 --- a/crates/rust-client/src/store/smt_forest.rs +++ b/crates/rust-client/src/store/smt_forest.rs @@ -76,6 +76,64 @@ impl AccountSmtForest { Ok(StorageMapWitness::new(proof, [key])?) } + // STAGED WITNESS QUERIES + // -------------------------------------------------------------------------------------------- + + /// Serves vault asset witnesses for `query_keys` against the vault root obtained by staging + /// the supplied asset updates onto `committed_vault_root`, *without* persisting the staged + /// tree. + /// + /// The staged tree is inserted into the forest (sharing nodes with the committed tree), + /// queried, then dropped: popping it decrements reference counts back to their committed + /// values, so the committed tree's nodes are left intact. The staged root is returned + /// alongside the witnesses so the caller can verify it matches the expected in-batch root. + /// + /// Unlike [`Self::get_asset_and_witness`], this also returns a witness for absent keys (an + /// emptiness proof), which the executor needs when an asset is being added to the vault. + pub fn staged_vault_asset_witnesses( + &mut self, + committed_vault_root: Word, + new_assets: impl Iterator, + removed_vault_keys: impl Iterator, + query_keys: impl IntoIterator, + ) -> Result<(Word, Vec), StoreError> { + let staged_root = + self.update_asset_nodes(committed_vault_root, new_assets, removed_vault_keys)?; + + let witnesses = query_keys + .into_iter() + .map(|key| { + let proof = self.forest.open(staged_root, key.into())?; + Ok(AssetWitness::new(proof)?) + }) + .collect::, StoreError>>(); + + // `update_asset_nodes` only adds a new root when the delta actually changed the tree; + // popping that staged root restores the committed tree's ref counts exactly. + if staged_root != committed_vault_root { + self.safe_pop_smts([staged_root]); + } + + witnesses.map(|witnesses| (staged_root, witnesses)) + } + + /// Serves a storage map witness for `query_key` against the map root obtained by staging the + /// supplied entries onto `committed_map_root`, *without* persisting the staged tree. See + /// [`Self::staged_vault_asset_witnesses`] for the staging / ref-count contract. + pub fn staged_storage_map_witness( + &mut self, + committed_map_root: Word, + entries: impl Iterator, + query_key: StorageMapKey, + ) -> Result<(Word, StorageMapWitness), StoreError> { + let staged_root = self.update_storage_map_nodes(committed_map_root, entries)?; + let witness = self.get_storage_map_item_witness(staged_root, query_key); + if staged_root != committed_map_root { + self.safe_pop_smts([staged_root]); + } + witness.map(|witness| (staged_root, witness)) + } + // ROOT LIFECYCLE // -------------------------------------------------------------------------------------------- diff --git a/crates/rust-client/src/transaction/batch/data_store.rs b/crates/rust-client/src/transaction/batch/data_store.rs index 279ffee5ea..8eac02a902 100644 --- a/crates/rust-client/src/transaction/batch/data_store.rs +++ b/crates/rust-client/src/transaction/batch/data_store.rs @@ -2,7 +2,13 @@ use alloc::collections::{BTreeMap, BTreeSet}; use alloc::sync::Arc; use alloc::vec::Vec; -use miden_protocol::account::{AccountId, PartialAccount, StorageMapKey, StorageMapWitness}; +use miden_protocol::account::{ + AccountDelta, + AccountId, + PartialAccount, + StorageMapKey, + StorageMapWitness, +}; use miden_protocol::asset::{AssetVaultKey, AssetWitness}; use miden_protocol::block::{BlockHeader, BlockNumber}; use miden_protocol::note::{NoteScript, NoteScriptRoot}; @@ -21,8 +27,13 @@ use crate::store::data_store::ClientDataStore; /// A [`DataStore`] that lets a [`crate::transaction::BatchBuilder`] stack in-memory account /// inputs for any number of local accounts. For each account pushed into the batch, a /// [`PartialAccount`] is cached; the executor sees the in-batch partial account state instead of -/// the stale store state. Witness reads for the cached account are resolved from the prior -/// transaction's execution advice, without reconstructing a full account. +/// the stale store state. +/// +/// Witness reads for the cached account are first resolved from the prior transaction's +/// execution advice (which covers the keys that transaction touched). Keys that no prior in-batch +/// transaction touched are absent from the advice; those are served by the [`crate::store::Store`] +/// by staging the accumulated in-batch delta onto its committed Merkle forest, so no full account +/// is ever reconstructed. pub(crate) struct InMemoryBatchDataStore { inner: ClientDataStore, current_accounts: BTreeMap, @@ -31,6 +42,9 @@ pub(crate) struct InMemoryBatchDataStore { struct CachedAccountState { account: PartialAccount, tx_inputs: TransactionInputs, + /// Accumulated delta from the account's committed state to the current in-batch state. Used + /// to serve witnesses for keys not present in `tx_inputs`' execution advice. + accumulated_delta: AccountDelta, } impl InMemoryBatchDataStore { @@ -40,10 +54,32 @@ impl InMemoryBatchDataStore { } /// Caches the post-transaction partial account and the transaction inputs carrying the - /// execution advice needed to serve later witness reads for this in-batch state. - pub(crate) fn cache_account(&mut self, account: PartialAccount, tx_inputs: TransactionInputs) { - self.current_accounts - .insert(account.id(), CachedAccountState { account, tx_inputs }); + /// execution advice for the just-executed transaction, and folds `delta` into the account's + /// accumulated in-batch delta so later transactions can resolve witnesses for any key. + pub(crate) fn cache_account( + &mut self, + account: PartialAccount, + tx_inputs: TransactionInputs, + delta: AccountDelta, + ) -> Result<(), ClientError> { + match self.current_accounts.get_mut(&account.id()) { + Some(state) => { + state.accumulated_delta.merge(delta)?; + state.account = account; + state.tx_inputs = tx_inputs; + }, + None => { + self.current_accounts.insert( + account.id(), + CachedAccountState { + account, + tx_inputs, + accumulated_delta: delta, + }, + ); + }, + } + Ok(()) } /// Returns the inner [`ClientDataStore`]'s MAST store so callers can load account @@ -105,25 +141,37 @@ impl DataStore for InMemoryBatchDataStore { vault_root: Word, vault_keys: BTreeSet, ) -> Result, DataStoreError> { - if let Some(state) = self.current_accounts.get(&account_id) { - let vault = state.account.vault(); - let in_batch_root = vault.root(); - if in_batch_root != vault_root { - return Err(DataStoreError::other(format!( - "vault root mismatch for account {account_id}: in-batch root = {in_batch_root:?}, requested root = {vault_root:?}", - ))); - } - state.tx_inputs.read_vault_asset_witnesses(vault_root, vault_keys).map_err(|err| { - DataStoreError::other_with_source( - format!( - "vault asset witness not found in in-batch account state for account {account_id}" - ), - err, - ) - }) - } else { - self.inner.get_vault_asset_witnesses(account_id, vault_root, vault_keys).await + let Some(state) = self.current_accounts.get(&account_id) else { + return self.inner.get_vault_asset_witnesses(account_id, vault_root, vault_keys).await; + }; + + let in_batch_root = state.account.vault().root(); + if in_batch_root != vault_root { + return Err(DataStoreError::other(format!( + "vault root mismatch for account {account_id}: in-batch root = {in_batch_root:?}, requested root = {vault_root:?}", + ))); } + + // Fast path: keys the prior in-batch transaction touched are in its execution advice. + if let Ok(witnesses) = + state.tx_inputs.read_vault_asset_witnesses(vault_root, vault_keys.clone()) + { + return Ok(witnesses); + } + + // Miss: a key no prior in-batch transaction touched. Serve it from the store by staging + // the accumulated in-batch delta onto the committed vault, without reconstructing the + // account. + self.inner + .store() + .vault_asset_witnesses_after_delta( + account_id, + state.accumulated_delta.clone(), + vault_root, + vault_keys, + ) + .await + .map_err(DataStoreError::from) } async fn get_storage_map_witness( @@ -132,23 +180,34 @@ impl DataStore for InMemoryBatchDataStore { map_root: Word, map_key: StorageMapKey, ) -> Result { - if let Some(state) = self.current_accounts.get(&account_id) { - if !state.account.storage().header().map_slot_roots().any(|root| root == map_root) { - return Err(DataStoreError::other(format!( - "storage map root not found in in-batch account state for account {account_id}: requested root = {map_root:?}", - ))); - } - state.tx_inputs.read_storage_map_witness(map_root, map_key).map_err(|err| { - DataStoreError::other_with_source( - format!( - "storage map witness not found in in-batch account state for account {account_id}" - ), - err, - ) - }) - } else { - self.inner.get_storage_map_witness(account_id, map_root, map_key).await + let Some(state) = self.current_accounts.get(&account_id) else { + return self.inner.get_storage_map_witness(account_id, map_root, map_key).await; + }; + + if !state.account.storage().header().map_slot_roots().any(|root| root == map_root) { + return Err(DataStoreError::other(format!( + "storage map root not found in in-batch account state for account {account_id}: requested root = {map_root:?}", + ))); + } + + // Fast path: a key the prior in-batch transaction touched is in its execution advice. + if let Ok(witness) = state.tx_inputs.read_storage_map_witness(map_root, map_key) { + return Ok(witness); } + + // Miss: a key no prior in-batch transaction touched. Serve it from the store by staging + // the accumulated in-batch delta onto the committed map, without reconstructing the + // account. + self.inner + .store() + .storage_map_witness_after_delta( + account_id, + state.accumulated_delta.clone(), + map_root, + map_key, + ) + .await + .map_err(DataStoreError::from) } async fn get_foreign_account_inputs( diff --git a/crates/rust-client/src/transaction/batch/mod.rs b/crates/rust-client/src/transaction/batch/mod.rs index 7a6732a085..4cb0ff4eee 100644 --- a/crates/rust-client/src/transaction/batch/mod.rs +++ b/crates/rust-client/src/transaction/batch/mod.rs @@ -327,7 +327,8 @@ where let current_account = partial_account_from_executed_transaction(&executed_transaction)?; let tx_inputs = executed_transaction.tx_inputs().clone(); - data_store.cache_account(current_account, tx_inputs); + let account_delta = executed_transaction.account_delta().clone(); + data_store.cache_account(current_account, tx_inputs, account_delta)?; validate_executed_transaction(&executed_transaction, &prep.output_recipients)?; TransactionResult::new(executed_transaction, prep.future_notes) diff --git a/crates/sqlite-store/src/account/accounts.rs b/crates/sqlite-store/src/account/accounts.rs index 37878354d6..9c62bfea0c 100644 --- a/crates/sqlite-store/src/account/accounts.rs +++ b/crates/sqlite-store/src/account/accounts.rs @@ -297,6 +297,53 @@ impl SqliteStore { Ok((item, witness)) } + /// Serves a storage map witness for `key` against the in-batch storage state obtained by + /// applying `delta` to the account's committed storage, *without* persisting the change. + /// + /// `map_root` is the in-batch root of the queried map slot. An unchanged slot's root is + /// already held by the forest and served directly; for a changed slot the slot's delta is + /// staged on its committed map root and the witness read at the resulting staged root. + pub(crate) fn storage_map_witness_after_delta( + conn: &mut Connection, + smt_forest: &Arc>, + account_id: AccountId, + delta: &AccountDelta, + map_root: Word, + key: StorageMapKey, + ) -> Result { + let committed_slots = query_storage_values(conn, account_id)?; + + let mut smt_forest = smt_forest + .write() + .map_err(|_| StoreError::DatabaseError("smt_forest write lock poisoned".to_string()))?; + + // Unchanged slot: the in-batch root equals the committed root the forest already holds, + // so the witness can be read directly. + if let Ok(witness) = smt_forest.get_storage_map_item_witness(map_root, key) { + return Ok(witness); + } + + // Changed slot: find the map slot whose committed root, after applying its delta, yields + // the requested in-batch root, and read the witness at the staged root. + for (slot_name, map_delta) in delta.storage().maps() { + let Some((slot_type, committed_root)) = committed_slots.get(slot_name) else { + continue; + }; + if *slot_type != StorageSlotType::Map { + continue; + } + + let entries = map_delta.entries().iter().map(|(map_key, value)| (*map_key, *value)); + let (staged_root, witness) = + smt_forest.staged_storage_map_witness(*committed_root, entries, key)?; + if staged_root == map_root { + return Ok(witness); + } + } + + Err(StoreError::AccountStorageRootNotFound(map_root)) + } + pub(crate) fn get_account_addresses( conn: &mut Connection, account_id: AccountId, diff --git a/crates/sqlite-store/src/account/vault.rs b/crates/sqlite-store/src/account/vault.rs index 9167528d01..675d236a1e 100644 --- a/crates/sqlite-store/src/account/vault.rs +++ b/crates/sqlite-store/src/account/vault.rs @@ -1,12 +1,13 @@ //! Vault/asset-related database operations for accounts. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::rc::Rc; +use std::sync::{Arc, RwLock}; use std::vec::Vec; use miden_client::Word; use miden_client::account::{AccountDelta, AccountHeader, AccountId}; -use miden_client::asset::{Asset, FungibleAsset, NonFungibleDeltaAction}; +use miden_client::asset::{Asset, AssetWitness, FungibleAsset, NonFungibleDeltaAction}; use miden_client::store::{AccountSmtForest, StoreError}; use miden_protocol::asset::AssetVaultKey; use miden_protocol::crypto::merkle::MerkleError; @@ -98,25 +99,60 @@ impl SqliteStore { account_id: AccountId, init_account_state: &AccountHeader, final_account_state: &AccountHeader, - mut updated_fungible_assets: BTreeMap, + updated_fungible_assets: BTreeMap, delta: &AccountDelta, ) -> Result<(), StoreError> { let nonce = final_account_state.nonce().as_canonical_u64(); let account_id_hex = account_id.to_hex(); let nonce_val = u64_to_value(nonce); - // Apply vault delta. This map will contain all updated assets (indexed by vault key), both - // fungible and non-fungible. + let (updated_assets_values, removed_vault_keys) = + Self::compute_vault_asset_updates(updated_fungible_assets, delta)?; + + Self::persist_vault_delta( + tx, + &account_id_hex, + &nonce_val, + &removed_vault_keys, + &updated_assets_values, + )?; + + let new_vault_root = smt_forest.update_asset_nodes( + init_account_state.vault_root(), + updated_assets_values.iter().copied(), + removed_vault_keys.iter().copied(), + )?; + if new_vault_root != final_account_state.vault_root() { + return Err(StoreError::MerkleStoreError(MerkleError::ConflictingRoots { + expected_root: final_account_state.vault_root(), + actual_root: new_vault_root, + })); + } + + Ok(()) + } + + /// Computes the absolute asset values and removed vault keys produced by applying `delta`'s + /// vault changes, given the account's current fungible balances for the affected keys. + /// + /// Fungible changes are signed amounts, so they are resolved against the supplied current + /// balances (`current_balance ± delta_amount`); non-fungible changes are absolute add/remove. + fn compute_vault_asset_updates( + mut current_fungible_assets: BTreeMap, + delta: &AccountDelta, + ) -> Result<(Vec, Vec), StoreError> { + // This map will contain all updated assets (indexed by vault key), both fungible and + // non-fungible. let mut updated_assets: BTreeMap = BTreeMap::new(); let mut removed_vault_keys: Vec = Vec::new(); - // We first process the fungible assets. Adding or subtracting them from the vault as - // requested. + // We first process the fungible assets, adding or subtracting them from the current + // balance as requested. for (vault_key, delta) in delta.vault().fungible().iter() { let delta_asset = FungibleAsset::new(vault_key.faucet_id(), delta.unsigned_abs())? .with_callbacks(vault_key.callback_flag()); - let asset = match updated_fungible_assets.remove(vault_key) { + let asset = match current_fungible_assets.remove(vault_key) { Some(asset) => { // If the asset exists, update it accordingly. if *delta >= 0 { @@ -154,28 +190,53 @@ impl SqliteStore { removed_vault_keys .extend(removed_nonfungible_assets.iter().map(|(asset, _)| asset.vault_key())); - let updated_assets_values: Vec = updated_assets.values().copied().collect(); - Self::persist_vault_delta( - tx, - &account_id_hex, - &nonce_val, - &removed_vault_keys, - &updated_assets_values, - )?; + Ok((updated_assets.into_values().collect(), removed_vault_keys)) + } - let new_vault_root = smt_forest.update_asset_nodes( - init_account_state.vault_root(), - updated_assets_values.iter().copied(), - removed_vault_keys.iter().copied(), + /// Serves vault asset witnesses for `vault_keys` against the in-batch vault state obtained by + /// applying `delta` to the account's committed vault, *without* persisting the change. + /// + /// The delta is staged on the in-memory [`AccountSmtForest`] (which already holds the + /// committed vault tree), the witnesses are read at the staged root, and the staged tree is + /// dropped. `vault_root` is the in-batch root the caller expects; a mismatch means the + /// supplied delta is inconsistent with the committed state. + pub(crate) fn vault_asset_witnesses_after_delta( + conn: &mut Connection, + smt_forest: &Arc>, + account_id: AccountId, + delta: &AccountDelta, + vault_root: Word, + vault_keys: BTreeSet, + ) -> Result, StoreError> { + let updated_fungible_assets = + Self::get_account_fungible_assets_for_delta(conn, account_id, delta)?; + let (updated_assets, removed_vault_keys) = + Self::compute_vault_asset_updates(updated_fungible_assets, delta)?; + + let committed_vault_root = Self::get_account_header(conn, account_id)? + .ok_or(StoreError::AccountDataNotFound(account_id))? + .0 + .vault_root(); + + let mut smt_forest = smt_forest + .write() + .map_err(|_| StoreError::DatabaseError("smt_forest write lock poisoned".to_string()))?; + + let (staged_root, witnesses) = smt_forest.staged_vault_asset_witnesses( + committed_vault_root, + updated_assets.into_iter(), + removed_vault_keys.into_iter(), + vault_keys, )?; - if new_vault_root != final_account_state.vault_root() { + + if staged_root != vault_root { return Err(StoreError::MerkleStoreError(MerkleError::ConflictingRoots { - expected_root: final_account_state.vault_root(), - actual_root: new_vault_root, + expected_root: vault_root, + actual_root: staged_root, })); } - Ok(()) + Ok(witnesses) } /// Persists vault delta changes: archives old values from latest to historical, diff --git a/crates/sqlite-store/src/lib.rs b/crates/sqlite-store/src/lib.rs index 9a3ca7eb30..1456c68110 100644 --- a/crates/sqlite-store/src/lib.rs +++ b/crates/sqlite-store/src/lib.rs @@ -23,6 +23,7 @@ use miden_client::Word; use miden_client::account::{ Account, AccountCode, + AccountDelta, AccountHeader, AccountId, AccountStorage, @@ -539,6 +540,48 @@ impl Store for SqliteStore { .await } + async fn vault_asset_witnesses_after_delta( + &self, + account_id: AccountId, + delta: AccountDelta, + vault_root: Word, + vault_keys: BTreeSet, + ) -> Result, StoreError> { + let smt_forest = self.smt_forest.clone(); + self.interact_with_connection(move |conn| { + SqliteStore::vault_asset_witnesses_after_delta( + conn, + &smt_forest, + account_id, + &delta, + vault_root, + vault_keys, + ) + }) + .await + } + + async fn storage_map_witness_after_delta( + &self, + account_id: AccountId, + delta: AccountDelta, + map_root: Word, + map_key: StorageMapKey, + ) -> Result { + let smt_forest = self.smt_forest.clone(); + self.interact_with_connection(move |conn| { + SqliteStore::storage_map_witness_after_delta( + conn, + &smt_forest, + account_id, + &delta, + map_root, + map_key, + ) + }) + .await + } + async fn get_addresses_by_account_id( &self, account_id: AccountId, diff --git a/crates/testing/miden-client-tests/src/tests/batch.rs b/crates/testing/miden-client-tests/src/tests/batch.rs index cbf0e829b9..84e2c3cd70 100644 --- a/crates/testing/miden-client-tests/src/tests/batch.rs +++ b/crates/testing/miden-client-tests/src/tests/batch.rs @@ -2,16 +2,19 @@ use std::collections::BTreeSet; use std::sync::Arc; use miden_client::account::AccountType; +use miden_client::assembly::CodeBuilder; use miden_client::asset::{Asset, FungibleAsset}; -use miden_client::auth::RPO_FALCON_SCHEME_ID; +use miden_client::auth::{AuthSchemeId, AuthSecretKey, AuthSingleSig, RPO_FALCON_SCHEME_ID}; use miden_client::builder::ClientBuilder; -use miden_client::keystore::FilesystemKeyStore; +use miden_client::keystore::{FilesystemKeyStore, Keystore}; use miden_client::note::{NoteType, NoteUpdateTracker}; use miden_client::rpc::NodeRpcClient; use miden_client::store::{StoreError, TransactionFilter}; use miden_client::testing::common::{ MINT_AMOUNT, + TRANSFER_AMOUNT, create_test_store_path, + insert_new_fungible_faucet, mint_and_consume, mint_note, setup_two_wallets_and_faucet, @@ -26,9 +29,21 @@ use miden_client::transaction::{ }; use miden_client::{ClientError, DebugMode}; use miden_client_sqlite_store::ClientBuilderSqliteExt; -use miden_protocol::Felt; +use miden_protocol::account::{ + AccountBuilder, + AccountComponent, + AccountComponentMetadata, + StorageMap, + StorageMapKey, + StorageSlot, + StorageSlotName, +}; use miden_protocol::crypto::rand::RandomCoin; +use miden_protocol::{Felt, Word}; +use miden_standards::account::AccountBuilderSchemaCommitmentExt; +use miden_standards::account::wallets::BasicWallet; use miden_testing::{Auth, MockChainBuilder, TxContextInput}; +use rand::RngCore; use crate::tests::create_test_client; @@ -330,6 +345,227 @@ async fn batch_builder_push_succeeds_when_balance_depends_on_prior_push() { assert!(block_num.as_u32() > 0); } +/// Storage map slot name used by the storage-map portion of the untouched-key regression test. +const BATCH_MAP_SLOT_NAME: &str = "miden::testing::batch_map::map"; +/// Two distinct keys of the same storage map; transaction 1 bumps A, transaction 2 bumps B. +const MAP_KEY_A: [Felt; 4] = [ + Felt::new_unchecked(11), + Felt::new_unchecked(11), + Felt::new_unchecked(11), + Felt::new_unchecked(11), +]; +const MAP_KEY_B: [Felt; 4] = [ + Felt::new_unchecked(22), + Felt::new_unchecked(22), + Felt::new_unchecked(22), + Felt::new_unchecked(22), +]; + +/// MASM for a procedure that increments the storage map entry at `key`. The body mirrors the +/// known-good `bump` sequence from the `storage_and_vault_proofs` test. +fn bump_proc(name: &str, key: &str) -> String { + format!( + r" +pub proc {name} + push.{key} + push.MAP_SLOT[0..2] + exec.::miden::protocol::active_account::get_map_item + add.1 + push.{key} + push.MAP_SLOT[0..2] + exec.::miden::protocol::native_account::set_map_item + dropw + dupw + push.MAP_SLOT[0..2] + exec.::miden::protocol::native_account::set_map_item + dropw dropw +end" + ) +} + +/// Account/script module exposing `bump_a` and `bump_b`, each mutating a distinct key of the same +/// storage map slot. The same module is installed on the account and linked into the tx scripts so +/// the `call` targets resolve to the account's procedures. +fn bump_map_module() -> String { + format!( + "use miden::core::word\n\nconst MAP_SLOT = word(\"{slot}\")\n{a}\n{b}", + slot = BATCH_MAP_SLOT_NAME, + a = bump_proc("bump_a", &Word::from(MAP_KEY_A).to_hex()), + b = bump_proc("bump_b", &Word::from(MAP_KEY_B).to_hex()), + ) +} + +/// A later transaction in a batch may touch a vault key *or* storage map key that an earlier +/// transaction in the same batch never touched. That key is absent from the earlier transaction's +/// execution advice, so the batch data store must serve its witness by staging the accumulated +/// in-batch delta onto the store's committed Merkle forest — not fail. Regression test for the +/// in-batch "untouched key" witness path, covering both the vault and storage-map cases. +/// +/// Vault case — `from` holds a balance of faucet G (committed); a note from a *different* faucet F +/// is left unconsumed: +/// - Push 1 consumes the F note → touches only the F vault key. +/// - Push 2 sends G to `to` → touches the G vault key, which push 1 never loaded. +/// +/// Storage-map case — a custom account has a map slot with two keys A and B: +/// - Push 1 bumps key A → touches only A. +/// - Push 2 bumps key B → touches B, which push 1 never loaded. +#[allow(clippy::too_many_lines)] +#[tokio::test] +async fn batch_builder_serves_witness_for_untouched_key() { + let (mut client, rpc_api, authenticator) = Box::pin(create_test_client()).await; + + let (from_account, to_account, consumed_faucet) = setup_two_wallets_and_faucet( + &mut client, + AccountType::Private, + &authenticator, + RPO_FALCON_SCHEME_ID, + ) + .await + .unwrap(); + + let from_id = from_account.id(); + let to_id = to_account.id(); + let consumed_faucet_id = consumed_faucet.id(); + + // A second, independent faucet whose balance `from` holds but never touches in push 1. + let (held_faucet, _) = insert_new_fungible_faucet( + &mut client, + AccountType::Private, + &authenticator, + RPO_FALCON_SCHEME_ID, + ) + .await + .unwrap(); + let held_faucet_id = held_faucet.id(); + client.sync_state().await.unwrap(); + + // Give `from` a committed balance of the held faucet. It is part of the committed vault but is + // NOT touched by the first in-batch transaction. + mint_and_consume(&mut client, from_id, held_faucet_id, NoteType::Private).await; + rpc_api.prove_block(); + client.sync_state().await.unwrap(); + + // Mint a note from the consumed faucet for `from`, left UNCONSUMED so push 1 can claim it. + let (_mint_tx_id, consumed_note) = + mint_note(&mut client, from_id, consumed_faucet_id, NoteType::Private).await; + rpc_api.prove_block(); + client.sync_state().await.unwrap(); + + // Push 1 consumes the note → touches only the consumed faucet's vault key. + let push1 = TransactionRequestBuilder::new() + .build_consume_notes(vec![consumed_note]) + .unwrap(); + + // Push 2 sends the held asset to `to` → touches the held faucet's vault key, absent from + // push 1's execution advice. + let held_asset = FungibleAsset::new(held_faucet_id, TRANSFER_AMOUNT).unwrap(); + let push2 = TransactionRequestBuilder::new() + .build_pay_to_id( + PaymentNoteDescription::new(vec![Asset::Fungible(held_asset)], from_id, to_id), + NoteType::Private, + client.rng(), + ) + .unwrap(); + + let block_num = Box::pin(async { + client + .new_transaction_batch() + .push(from_id, push1) + .await? + .push(from_id, push2) + .await? + .submit() + .await + }) + .await + .expect("submit should succeed: the untouched G vault key is served via the store forest"); + + assert!(block_num.as_u32() > 0); + + // --------------------------------------------------------------------------------------------- + // Storage-map case: a custom account with a map slot holding keys A and B. Push 1 bumps A, + // push 2 bumps B — B is never touched by push 1 and so is absent from its execution advice. + // --------------------------------------------------------------------------------------------- + let module = bump_map_module(); + + let component_code = CodeBuilder::default() + .compile_component_code("miden::testing::batch_map_component", module.clone()) + .unwrap(); + + let mut storage_map = StorageMap::new(); + let initial_value: Word = + [Felt::from(0u32), Felt::from(0u32), Felt::from(0u32), Felt::from(1u32)].into(); + storage_map.insert(StorageMapKey::new(MAP_KEY_A.into()), initial_value).unwrap(); + storage_map.insert(StorageMapKey::new(MAP_KEY_B.into()), initial_value).unwrap(); + + let map_slot = + StorageSlot::with_map(StorageSlotName::new(BATCH_MAP_SLOT_NAME).unwrap(), storage_map); + let map_component = AccountComponent::new( + component_code, + vec![map_slot], + AccountComponentMetadata::new("miden::testing::batch_map_component"), + ) + .unwrap(); + + let key_pair = AuthSecretKey::new_falcon512_poseidon2(); + let pub_key = key_pair.public_key(); + let mut init_seed = [0u8; 32]; + client.rng().fill_bytes(&mut init_seed); + let map_account = AccountBuilder::new(init_seed) + .account_type(AccountType::Public) + .with_auth_component(AuthSingleSig::new( + pub_key.to_commitment(), + AuthSchemeId::Falcon512Poseidon2, + )) + .with_component(BasicWallet) + .with_component(map_component) + .build_with_schema_commitment() + .unwrap(); + let map_account_id = map_account.id(); + authenticator.add_key(&key_pair, map_account_id).await.unwrap(); + client.add_account(&map_account, false).await.unwrap(); + + // Commit the account on chain so the batch runs against committed state (as the wallets above). + mint_and_consume(&mut client, map_account_id, held_faucet_id, NoteType::Private).await; + rpc_api.prove_block(); + client.sync_state().await.unwrap(); + + // The same module is linked into both scripts so `call.batch_map::bump_*` resolves to the + // account's procedures. + let script_a = CodeBuilder::new() + .with_linked_module("external_contract::batch_map", module.clone()) + .unwrap() + .compile_tx_script( + "use external_contract::batch_map\nbegin\n call.batch_map::bump_a\nend", + ) + .unwrap(); + let script_b = CodeBuilder::new() + .with_linked_module("external_contract::batch_map", module.clone()) + .unwrap() + .compile_tx_script( + "use external_contract::batch_map\nbegin\n call.batch_map::bump_b\nend", + ) + .unwrap(); + + let map_push1 = TransactionRequestBuilder::new().custom_script(script_a).build().unwrap(); + let map_push2 = TransactionRequestBuilder::new().custom_script(script_b).build().unwrap(); + + let map_block_num = Box::pin(async { + client + .new_transaction_batch() + .push(map_account_id, map_push1) + .await? + .push(map_account_id, map_push2) + .await? + .submit() + .await + }) + .await + .expect("submit should succeed: the untouched map key B is served via the store forest"); + + assert!(map_block_num.as_u32() > 0); +} + /// Verify that submitting an empty batch (no pushes) returns `BatchBuilderError::Empty`. #[tokio::test] async fn batch_builder_empty_submit_returns_empty_error() { From 62d407f24cdf344cf7e83a64e79ad6df1dbdaa16 Mon Sep 17 00:00:00 2001 From: Ignacio Amigo Date: Thu, 25 Jun 2026 22:45:04 -0300 Subject: [PATCH 6/6] test: simplify test --- .../miden-client-tests/src/tests/batch.rs | 178 ++---------------- 1 file changed, 14 insertions(+), 164 deletions(-) diff --git a/crates/testing/miden-client-tests/src/tests/batch.rs b/crates/testing/miden-client-tests/src/tests/batch.rs index 84e2c3cd70..30fc0b7dec 100644 --- a/crates/testing/miden-client-tests/src/tests/batch.rs +++ b/crates/testing/miden-client-tests/src/tests/batch.rs @@ -2,11 +2,10 @@ use std::collections::BTreeSet; use std::sync::Arc; use miden_client::account::AccountType; -use miden_client::assembly::CodeBuilder; use miden_client::asset::{Asset, FungibleAsset}; -use miden_client::auth::{AuthSchemeId, AuthSecretKey, AuthSingleSig, RPO_FALCON_SCHEME_ID}; +use miden_client::auth::RPO_FALCON_SCHEME_ID; use miden_client::builder::ClientBuilder; -use miden_client::keystore::{FilesystemKeyStore, Keystore}; +use miden_client::keystore::FilesystemKeyStore; use miden_client::note::{NoteType, NoteUpdateTracker}; use miden_client::rpc::NodeRpcClient; use miden_client::store::{StoreError, TransactionFilter}; @@ -29,21 +28,9 @@ use miden_client::transaction::{ }; use miden_client::{ClientError, DebugMode}; use miden_client_sqlite_store::ClientBuilderSqliteExt; -use miden_protocol::account::{ - AccountBuilder, - AccountComponent, - AccountComponentMetadata, - StorageMap, - StorageMapKey, - StorageSlot, - StorageSlotName, -}; +use miden_protocol::Felt; use miden_protocol::crypto::rand::RandomCoin; -use miden_protocol::{Felt, Word}; -use miden_standards::account::AccountBuilderSchemaCommitmentExt; -use miden_standards::account::wallets::BasicWallet; use miden_testing::{Auth, MockChainBuilder, TxContextInput}; -use rand::RngCore; use crate::tests::create_test_client; @@ -345,73 +332,19 @@ async fn batch_builder_push_succeeds_when_balance_depends_on_prior_push() { assert!(block_num.as_u32() > 0); } -/// Storage map slot name used by the storage-map portion of the untouched-key regression test. -const BATCH_MAP_SLOT_NAME: &str = "miden::testing::batch_map::map"; -/// Two distinct keys of the same storage map; transaction 1 bumps A, transaction 2 bumps B. -const MAP_KEY_A: [Felt; 4] = [ - Felt::new_unchecked(11), - Felt::new_unchecked(11), - Felt::new_unchecked(11), - Felt::new_unchecked(11), -]; -const MAP_KEY_B: [Felt; 4] = [ - Felt::new_unchecked(22), - Felt::new_unchecked(22), - Felt::new_unchecked(22), - Felt::new_unchecked(22), -]; - -/// MASM for a procedure that increments the storage map entry at `key`. The body mirrors the -/// known-good `bump` sequence from the `storage_and_vault_proofs` test. -fn bump_proc(name: &str, key: &str) -> String { - format!( - r" -pub proc {name} - push.{key} - push.MAP_SLOT[0..2] - exec.::miden::protocol::active_account::get_map_item - add.1 - push.{key} - push.MAP_SLOT[0..2] - exec.::miden::protocol::native_account::set_map_item - dropw - dupw - push.MAP_SLOT[0..2] - exec.::miden::protocol::native_account::set_map_item - dropw dropw -end" - ) -} - -/// Account/script module exposing `bump_a` and `bump_b`, each mutating a distinct key of the same -/// storage map slot. The same module is installed on the account and linked into the tx scripts so -/// the `call` targets resolve to the account's procedures. -fn bump_map_module() -> String { - format!( - "use miden::core::word\n\nconst MAP_SLOT = word(\"{slot}\")\n{a}\n{b}", - slot = BATCH_MAP_SLOT_NAME, - a = bump_proc("bump_a", &Word::from(MAP_KEY_A).to_hex()), - b = bump_proc("bump_b", &Word::from(MAP_KEY_B).to_hex()), - ) -} - -/// A later transaction in a batch may touch a vault key *or* storage map key that an earlier -/// transaction in the same batch never touched. That key is absent from the earlier transaction's -/// execution advice, so the batch data store must serve its witness by staging the accumulated -/// in-batch delta onto the store's committed Merkle forest — not fail. Regression test for the -/// in-batch "untouched key" witness path, covering both the vault and storage-map cases. -/// -/// Vault case — `from` holds a balance of faucet G (committed); a note from a *different* faucet F -/// is left unconsumed: -/// - Push 1 consumes the F note → touches only the F vault key. -/// - Push 2 sends G to `to` → touches the G vault key, which push 1 never loaded. +/// A later transaction in a batch may touch a vault key that an earlier transaction in the same +/// batch never touched. That key is absent from the earlier transaction's execution advice, so the +/// batch data store must serve its witness by staging the accumulated in-batch delta onto the +/// store's committed Merkle forest — not fail. Regression test for the in-batch "untouched key" +/// witness path. /// -/// Storage-map case — a custom account has a map slot with two keys A and B: -/// - Push 1 bumps key A → touches only A. -/// - Push 2 bumps key B → touches B, which push 1 never loaded. -#[allow(clippy::too_many_lines)] +/// Setup: `from` holds a balance of a "held" faucet (committed); a note from a *different* +/// "consumed" faucet is left unconsumed. +/// - Push 1 consumes the note → touches only the consumed faucet's vault key. +/// - Push 2 sends the held asset to `to` → touches the held faucet's vault key, which push 1 never +/// loaded. #[tokio::test] -async fn batch_builder_serves_witness_for_untouched_key() { +async fn batch_builder_serves_witness_for_untouched_vault_key() { let (mut client, rpc_api, authenticator) = Box::pin(create_test_client()).await; let (from_account, to_account, consumed_faucet) = setup_two_wallets_and_faucet( @@ -481,89 +414,6 @@ async fn batch_builder_serves_witness_for_untouched_key() { .expect("submit should succeed: the untouched G vault key is served via the store forest"); assert!(block_num.as_u32() > 0); - - // --------------------------------------------------------------------------------------------- - // Storage-map case: a custom account with a map slot holding keys A and B. Push 1 bumps A, - // push 2 bumps B — B is never touched by push 1 and so is absent from its execution advice. - // --------------------------------------------------------------------------------------------- - let module = bump_map_module(); - - let component_code = CodeBuilder::default() - .compile_component_code("miden::testing::batch_map_component", module.clone()) - .unwrap(); - - let mut storage_map = StorageMap::new(); - let initial_value: Word = - [Felt::from(0u32), Felt::from(0u32), Felt::from(0u32), Felt::from(1u32)].into(); - storage_map.insert(StorageMapKey::new(MAP_KEY_A.into()), initial_value).unwrap(); - storage_map.insert(StorageMapKey::new(MAP_KEY_B.into()), initial_value).unwrap(); - - let map_slot = - StorageSlot::with_map(StorageSlotName::new(BATCH_MAP_SLOT_NAME).unwrap(), storage_map); - let map_component = AccountComponent::new( - component_code, - vec![map_slot], - AccountComponentMetadata::new("miden::testing::batch_map_component"), - ) - .unwrap(); - - let key_pair = AuthSecretKey::new_falcon512_poseidon2(); - let pub_key = key_pair.public_key(); - let mut init_seed = [0u8; 32]; - client.rng().fill_bytes(&mut init_seed); - let map_account = AccountBuilder::new(init_seed) - .account_type(AccountType::Public) - .with_auth_component(AuthSingleSig::new( - pub_key.to_commitment(), - AuthSchemeId::Falcon512Poseidon2, - )) - .with_component(BasicWallet) - .with_component(map_component) - .build_with_schema_commitment() - .unwrap(); - let map_account_id = map_account.id(); - authenticator.add_key(&key_pair, map_account_id).await.unwrap(); - client.add_account(&map_account, false).await.unwrap(); - - // Commit the account on chain so the batch runs against committed state (as the wallets above). - mint_and_consume(&mut client, map_account_id, held_faucet_id, NoteType::Private).await; - rpc_api.prove_block(); - client.sync_state().await.unwrap(); - - // The same module is linked into both scripts so `call.batch_map::bump_*` resolves to the - // account's procedures. - let script_a = CodeBuilder::new() - .with_linked_module("external_contract::batch_map", module.clone()) - .unwrap() - .compile_tx_script( - "use external_contract::batch_map\nbegin\n call.batch_map::bump_a\nend", - ) - .unwrap(); - let script_b = CodeBuilder::new() - .with_linked_module("external_contract::batch_map", module.clone()) - .unwrap() - .compile_tx_script( - "use external_contract::batch_map\nbegin\n call.batch_map::bump_b\nend", - ) - .unwrap(); - - let map_push1 = TransactionRequestBuilder::new().custom_script(script_a).build().unwrap(); - let map_push2 = TransactionRequestBuilder::new().custom_script(script_b).build().unwrap(); - - let map_block_num = Box::pin(async { - client - .new_transaction_batch() - .push(map_account_id, map_push1) - .await? - .push(map_account_id, map_push2) - .await? - .submit() - .await - }) - .await - .expect("submit should succeed: the untouched map key B is served via the store forest"); - - assert!(map_block_num.as_u32() > 0); } /// Verify that submitting an empty batch (no pushes) returns `BatchBuilderError::Empty`.