diff --git a/crates/rust-client/src/account/account_reader.rs b/crates/rust-client/src/account/account_reader.rs index 3e57490ea..855957e92 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/errors.rs b/crates/rust-client/src/errors.rs index 21a4c5a92..ffe91a9e5 100644 --- a/crates/rust-client/src/errors.rs +++ b/crates/rust-client/src/errors.rs @@ -6,7 +6,14 @@ 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, + AccountPatchError, + AssetError, + NetworkIdError, +}; use miden_protocol::errors::{ NoteError, PartialBlockchainError, @@ -78,6 +85,10 @@ pub enum ClientError { AccountAlreadyTracked(AccountId), #[error("account error")] AccountError(#[from] AccountError), + #[error("account delta error")] + AccountDeltaError(#[from] AccountDeltaError), + #[error("account patch error")] + AccountPatchError(#[from] AccountPatchError), #[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 7e56f28d9..f1073e6e6 100644 --- a/crates/rust-client/src/store/data_store/mod.rs +++ b/crates/rust-client/src/store/data_store/mod.rs @@ -69,6 +69,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 a19912191..18ae79b68 100644 --- a/crates/rust-client/src/store/errors.rs +++ b/crates/rust-client/src/store/errors.rs @@ -98,6 +98,8 @@ pub enum StoreError { VaultKeyNotTracked(AssetId, 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 b93e09518..1e0067235 100644 --- a/crates/rust-client/src/store/mod.rs +++ b/crates/rust-client/src/store/mod.rs @@ -31,6 +31,7 @@ use miden_protocol::account::{ AccountCode, AccountHeader, AccountId, + AccountPatch, AccountStorage, StorageMapKey, StorageMapWitness, @@ -668,6 +669,48 @@ pub trait Store: Send + Sync { } } + // IN-BATCH (STAGED) WITNESSES + // -------------------------------------------------------------------------------------------- + + /// Returns vault asset witnesses for `asset_ids` as the account's vault would look after + /// applying `patch` 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 patch on + /// that forest. + async fn vault_asset_witnesses_after_patch( + &self, + account_id: AccountId, + patch: AccountPatch, + vault_root: Word, + asset_ids: BTreeSet, + ) -> Result, StoreError> { + let _ = (account_id, patch, vault_root, asset_ids); + Err(StoreError::UnsupportedOperation("vault_asset_witnesses_after_patch")) + } + + /// Returns the storage map witness for `map_key` as the account's storage would look after + /// applying `patch` 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_patch`] for the rationale and the default + /// behavior. + async fn storage_map_witness_after_patch( + &self, + account_id: AccountId, + patch: AccountPatch, + map_root: Word, + map_key: StorageMapKey, + ) -> Result { + let _ = (account_id, patch, map_root, map_key); + Err(StoreError::UnsupportedOperation("storage_map_witness_after_patch")) + } + // PARTIAL ACCOUNTS // -------------------------------------------------------------------------------------------- diff --git a/crates/rust-client/src/store/smt_forest.rs b/crates/rust-client/src/store/smt_forest.rs index 99eb97f8e..ef53ce766 100644 --- a/crates/rust-client/src/store/smt_forest.rs +++ b/crates/rust-client/src/store/smt_forest.rs @@ -77,6 +77,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(|asset_id| { + let proof = self.forest.open(staged_root, asset_id.hash().into())?; + Ok(AssetWitness::new(proof, [asset_id])?) + }) + .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 d5f65dad9..55c32054e 100644 --- a/crates/rust-client/src/transaction/batch/data_store.rs +++ b/crates/rust-client/src/transaction/batch/data_store.rs @@ -4,17 +4,16 @@ use alloc::vec::Vec; use miden_protocol::Word; use miden_protocol::account::{ - Account, AccountId, + AccountPatch, PartialAccount, StorageMapKey, StorageMapWitness, - StorageSlotContent, }; use miden_protocol::asset::{AssetId, AssetWitness}; use miden_protocol::block::{BlockHeader, BlockNumber}; 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_tx::{ DataStore, @@ -24,18 +23,35 @@ use miden_tx::{ TransactionMastStore, }; +use crate::ClientError; +use crate::account::AccountReader; use crate::store::data_store::ClientDataStore; // IN-MEMORY BATCH DATA STORE // ================================================================================================ /// 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 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 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, + current_accounts: BTreeMap, +} + +struct CachedAccountState { + account: PartialAccount, + tx_inputs: TransactionInputs, + /// Accumulated patch from the account's committed state to the current in-batch state. Its + /// absolute values are staged onto the committed Merkle forest to serve witnesses for keys + /// not present in `tx_inputs`' execution advice. + accumulated_patch: AccountPatch, } impl InMemoryBatchDataStore { @@ -44,18 +60,33 @@ impl InMemoryBatchDataStore { Self { inner, current_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. - pub(crate) fn get_account(&self, id: AccountId) -> Option<&Account> { - self.current_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); + /// Caches the post-transaction partial account and the transaction inputs carrying the + /// execution advice for the just-executed transaction, and folds `patch` into the account's + /// accumulated in-batch patch so later transactions can resolve witnesses for any key. + pub(crate) fn cache_account( + &mut self, + account: PartialAccount, + tx_inputs: TransactionInputs, + patch: AccountPatch, + ) -> Result<(), ClientError> { + match self.current_accounts.get_mut(&account.id()) { + Some(state) => { + state.accumulated_patch.merge(patch)?; + state.account = account; + state.tx_inputs = tx_inputs; + }, + None => { + self.current_accounts.insert( + account.id(), + CachedAccountState { + account, + tx_inputs, + accumulated_patch: patch, + }, + ); + }, + } + Ok(()) } /// Returns the inner [`ClientDataStore`]'s MAST store so callers can load account @@ -78,6 +109,18 @@ impl InMemoryBatchDataStore { pub(crate) fn register_note_scripts(&self, note_scripts: impl IntoIterator) { self.inner.register_note_scripts(note_scripts); } + + pub(crate) async fn current_account( + &self, + 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()); + } + + account_reader.partial_account().await + } } // DATA STORE IMPL @@ -92,8 +135,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 = PartialAccount::from(account); + if let Some(state) = self.current_accounts.get(&account_id) { + partial_account = state.account.clone(); } Ok((partial_account, block_header, partial_blockchain)) @@ -105,19 +148,37 @@ impl DataStore for InMemoryBatchDataStore { vault_root: Word, asset_ids: BTreeSet, ) -> Result, DataStoreError> { - if let Some(account) = self.current_accounts.get(&account_id) { - let vault = 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 = asset_ids.into_iter().map(|key| vault.open(key)).collect(); - Ok(witnesses) - } else { - self.inner.get_vault_asset_witnesses(account_id, vault_root, asset_ids).await + let Some(state) = self.current_accounts.get(&account_id) else { + return self.inner.get_vault_asset_witnesses(account_id, vault_root, asset_ids).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, asset_ids.clone()) + { + return Ok(witnesses); } + + // Miss: a key no prior in-batch transaction touched. Serve it from the store by staging + // the accumulated in-batch patch onto the committed vault, without reconstructing the + // account. + self.inner + .store() + .vault_asset_witnesses_after_patch( + account_id, + state.accumulated_patch.clone(), + vault_root, + asset_ids, + ) + .await + .map_err(DataStoreError::from) } async fn get_storage_map_witness( @@ -126,19 +187,34 @@ impl DataStore for InMemoryBatchDataStore { map_root: Word, 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)); - } - } + 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:?}", ))); } - self.inner.get_storage_map_witness(account_id, map_root, map_key).await + + // 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 patch onto the committed map, without reconstructing the + // account. + self.inner + .store() + .storage_map_witness_after_patch( + account_id, + state.accumulated_patch.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 1c03bdc71..4eccd4145 100644 --- a/crates/rust-client/src/transaction/batch/mod.rs +++ b/crates/rust-client/src/transaction/batch/mod.rs @@ -55,12 +55,31 @@ 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, + AccountStorageHeader, + PartialAccount, + PartialStorage, + PartialStorageMap, + StorageMap, + StorageMapPatch, + StorageSlotHeader, + StorageSlotPatch, + 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_protocol::transaction::{ + ExecutedTransaction, + PartialBlockchain, + ProvenTransaction, + TransactionInputs, +}; +use miden_protocol::{MIN_PROOF_SECURITY_LEVEL, ZERO}; +use miden_tx::DataStoreError; use miden_tx::auth::TransactionAuthenticator; use miden_tx_batch::{BatchExecutor, LocalBatchProver}; @@ -262,8 +281,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, @@ -273,23 +292,15 @@ 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 - }; + 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_reader).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 { @@ -302,7 +313,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?; } @@ -311,16 +327,153 @@ where .execute_transaction(account_id, prep.block_num, notes, prep.tx_args) .await?; - // Cache new account state in memory data store. + // Cache the post-transaction in-batch state: the rebuilt partial account (headers only), the + // execution advice for the fast-path witness lookups, and the account patch whose absolute + // values are staged onto the committed forest to serve witnesses for untouched keys. + let current_account = partial_account_from_executed_transaction(&executed_transaction)?; + let tx_inputs = executed_transaction.tx_inputs().clone(); + let patch = executed_transaction.account_patch().clone(); + data_store.cache_account(current_account, tx_inputs, patch)?; + + 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 patch = executed_transaction.account_patch(); - let account = if patch.is_full_state() { - Account::try_from(patch).map_err(ClientError::AccountError)? + let code = patch.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_patch(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 { - account.apply_patch(patch).map_err(ClientError::AccountError)?; - account + None }; - data_store.cache_account(account_id, account); - validate_executed_transaction(&executed_transaction, &prep.output_recipients)?; - TransactionResult::new(executed_transaction, prep.future_notes) + 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_patch( + executed_transaction: &ExecutedTransaction, +) -> Result { + let initial_account = executed_transaction.initial_account(); + let final_account = executed_transaction.final_account(); + let storage_patch = executed_transaction.account_patch().storage(); + + let mut slots = Vec::new(); + for slot in initial_account.storage().header().slots() { + let new_slot_value = match storage_patch.get(slot.name()) { + None => slot.value(), + Some(StorageSlotPatch::Value(value_patch)) => { + 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() + ))); + } + // A removed value slot commits to the empty word. + value_patch.value().unwrap_or_default() + }, + Some(StorageSlotPatch::Map(map_patch)) => { + 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_patch)? + }, + }; + + 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_patch: &StorageMapPatch, +) -> Result { + // A removed map slot reduces to the empty map root. + let Some(entries) = map_patch.entries() else { + return Ok(StorageMap::default().root()); + }; + let entries = entries.as_map(); + + // The patch carries absolute new values, so each changed key is inserted verbatim onto a + // partial SMT seeded with the initial map witnesses from this transaction's execution advice. + let mut partial_smt = PartialSmt::new(initial_root); + + for map_key in 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 entries { + partial_smt.insert(map_key.hash().as_word(), *value).map_err(|err| { + DataStoreError::other_with_source( + "failed to apply storage map patch 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 048e7a330..22898b567 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; @@ -352,24 +352,42 @@ 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 { - let account_id = account.id(); + self.prepare_transaction_inner(account.id(), transaction_request, Some(account)) + .await + } + + pub(crate) async fn prepare_transaction_for_batch( + &self, + account: &PartialAccount, + transaction_request: TransactionRequest, + ) -> Result { + self.prepare_transaction_inner(account.id(), transaction_request, None).await + } + + async fn prepare_transaction_inner( + &self, + account_id: AccountId, + transaction_request: TransactionRequest, + account_to_validate: Option<&Account>, + ) -> Result { self.validate_recency().await?; - validate_account_request(&transaction_request, account)?; + if let Some(account) = account_to_validate { + validate_account_request(&transaction_request, account)?; + } // Retrieve all input notes from the store. let mut stored_note_records = self @@ -693,18 +711,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()); + + self.get_valid_input_notes_with_data_store(&data_store, account.id(), input_notes, tx_args) + .await + } - data_store.mast_store().load_account_code(account.code()); - let execution = NoteConsumptionChecker::new(&self.build_executor(&data_store)?) + 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(), diff --git a/crates/sqlite-store/src/account/accounts.rs b/crates/sqlite-store/src/account/accounts.rs index f9f9b9457..911588997 100644 --- a/crates/sqlite-store/src/account/accounts.rs +++ b/crates/sqlite-store/src/account/accounts.rs @@ -292,6 +292,57 @@ impl SqliteStore { Ok((item, witness)) } + /// Serves a storage map witness for `key` against the in-batch storage state obtained by + /// applying `patch` 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 patch is + /// staged on its committed map root and the witness read at the resulting staged root. + pub(crate) fn storage_map_witness_after_patch( + conn: &mut Connection, + smt_forest: &Arc>, + account_id: AccountId, + patch: &AccountPatch, + 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 patch, yields + // the requested in-batch root, and read the witness at the staged root. The patch's + // absolute entries are staged verbatim. + for (slot_name, map_patch) in patch.storage().maps() { + let Some((slot_type, committed_root)) = committed_slots.get(slot_name) else { + continue; + }; + if *slot_type != StorageSlotType::Map { + continue; + } + let Some(entries) = map_patch.entries() else { + continue; + }; + + let entries = entries.as_map().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 9224047d4..d4b91f87d 100644 --- a/crates/sqlite-store/src/account/vault.rs +++ b/crates/sqlite-store/src/account/vault.rs @@ -1,16 +1,18 @@ //! Vault/asset-related database operations for accounts. +use std::collections::BTreeSet; use std::rc::Rc; +use std::sync::{Arc, RwLock}; use std::vec::Vec; -use miden_client::Serializable; -use miden_client::account::{AccountHeader, AccountId, AccountVaultPatch}; -use miden_client::asset::Asset; +use miden_client::account::{AccountHeader, AccountId, AccountPatch, AccountVaultPatch}; +use miden_client::asset::{Asset, AssetWitness}; use miden_client::store::{AccountSmtForest, StoreError}; +use miden_client::{Serializable, Word}; use miden_protocol::asset::AssetId; use miden_protocol::crypto::merkle::MerkleError; use rusqlite::types::Value; -use rusqlite::{OptionalExtension, Transaction, params}; +use rusqlite::{Connection, OptionalExtension, Transaction, params}; use crate::sql_error::SqlResultExt; use crate::{SqliteStore, insert_sql, subst, u64_to_value}; @@ -48,7 +50,7 @@ impl SqliteStore { Ok(()) } - /// Applies vault delta changes to the account state, updating fungible and non-fungible assets. + /// Applies the vault patch to the account state, updating fungible and non-fungible assets. /// /// The function updates the SMT forest with all asset changes and verifies that the resulting /// vault root matches the expected final state. It archives old values from latest to @@ -95,6 +97,50 @@ impl SqliteStore { Ok(()) } + /// Serves vault asset witnesses for `asset_ids` against the in-batch vault state obtained by + /// applying `patch` to the account's committed vault, *without* persisting the change. + /// + /// The patch's absolute asset values are 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 patch is inconsistent with the committed state. + pub(crate) fn vault_asset_witnesses_after_patch( + conn: &mut Connection, + smt_forest: &Arc>, + account_id: AccountId, + patch: &AccountPatch, + vault_root: Word, + asset_ids: BTreeSet, + ) -> Result, StoreError> { + let updated_assets: Vec = patch.vault().updated_assets().collect(); + let removed_vault_ids: Vec = patch.vault().removed_asset_ids().copied().collect(); + + 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_ids.into_iter(), + asset_ids, + )?; + + if staged_root != vault_root { + return Err(StoreError::MerkleStoreError(MerkleError::ConflictingRoots { + expected_root: vault_root, + actual_root: staged_root, + })); + } + + Ok(witnesses) + } + /// Persists vault delta changes: archives old values from latest to historical, /// then updates latest (deletes removed assets, inserts/updates changed assets). fn persist_vault_delta( diff --git a/crates/sqlite-store/src/lib.rs b/crates/sqlite-store/src/lib.rs index 0177d6006..08bb023e2 100644 --- a/crates/sqlite-store/src/lib.rs +++ b/crates/sqlite-store/src/lib.rs @@ -25,6 +25,7 @@ use miden_client::account::{ AccountCode, AccountHeader, AccountId, + AccountPatch, AccountStorage, Address, StorageMapKey, @@ -530,6 +531,48 @@ impl Store for SqliteStore { .await } + async fn vault_asset_witnesses_after_patch( + &self, + account_id: AccountId, + patch: AccountPatch, + vault_root: Word, + asset_ids: BTreeSet, + ) -> Result, StoreError> { + let smt_forest = self.smt_forest.clone(); + self.interact_with_connection(move |conn| { + SqliteStore::vault_asset_witnesses_after_patch( + conn, + &smt_forest, + account_id, + &patch, + vault_root, + asset_ids, + ) + }) + .await + } + + async fn storage_map_witness_after_patch( + &self, + account_id: AccountId, + patch: AccountPatch, + 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_patch( + conn, + &smt_forest, + account_id, + &patch, + 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 c7f94d171..a0613e083 100644 --- a/crates/testing/miden-client-tests/src/tests/batch.rs +++ b/crates/testing/miden-client-tests/src/tests/batch.rs @@ -12,7 +12,9 @@ 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, @@ -329,6 +331,90 @@ async fn batch_builder_push_succeeds_when_balance_depends_on_prior_push() { assert!(block_num.as_u32() > 0); } +/// 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. +/// +/// 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_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( + &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); +} + /// Verify that submitting an empty batch (no pushes) returns `BatchBuilderError::Empty`. #[tokio::test] async fn batch_builder_empty_submit_returns_empty_error() {