refactor(batch): keep PartialAccount instead of Account

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<PartialAccount, ClientError> {
+        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<Vec<Address>, ClientError> {
         self.store

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(

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<dyn Store> {
+        &self.store
+    }
+
     /// Stores the provided foreign account inputs so they can be served to the executor upon
     /// request.
     pub fn register_foreign_account_inputs(

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<StoreError> for DataStoreError {

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<AssetVaultKey>,
+    ) -> Result<Vec<AssetWitness>, 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<StorageMapWitness, StoreError> {
+        let _ = (account_id, delta, map_root, map_key);
+        Err(StoreError::UnsupportedOperation("storage_map_witness_after_delta"))
+    }
+
     // PARTIAL ACCOUNTS
     // --------------------------------------------------------------------------------------------
 

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<Item = Asset>,
+        removed_vault_keys: impl Iterator<Item = AssetVaultKey>,
+        query_keys: impl IntoIterator<Item = AssetVaultKey>,
+    ) -> Result<(Word, Vec<AssetWitness>), 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::<Result<Vec<_>, 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<Item = (StorageMapKey, Word)>,
+        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
     // --------------------------------------------------------------------------------------------