feat(store): cache account witnesses for tracked accounts

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Breaking Changes
 
+* [BREAKING][param][rust] The `Store` trait requires five account-witness registry methods: `track_account_witness`, `untrack_account_witness`, `tracked_account_witnesses`, `get_account_witness` and `update_account_witness`. They have no default bodies, so every out-of-tree implementation must provide them ([#2476](https://github.com/0xMiden/rust-sdk/pull/2476)).
 * [BREAKING][removal][rust] Removed `Client::try_get_account`. Use `Client::get_account` and handle the `None` case, or `Client::account_reader` for existence checks and single-field reads that don't need the full materialized account ([#2362](https://github.com/0xMiden/rust-sdk/pull/2362)).
 * [BREAKING][behavior][rpc] The `GetAccount` response no longer carries one SMT opening per requested storage map key. A slot queried with specific keys now comes back as a single partial SMT covering all of them, alongside the original unhashed keys, so the client requires a node that speaks this format ([#2362](https://github.com/0xMiden/rust-sdk/pull/2362)).
 * [BREAKING][type][rust] `StorageMapEntries::EntriesWithProofs(Vec<SmtProof>)` is replaced by `StorageMapEntries::PartialMap { map_keys, partial_smt }`, which carries the values only inside the tree: read one by hashing its raw key and calling `PartialSmt::get_value`. The enum also gained a `LimitExceeded` variant and `AccountStorageMapDetails::too_many_entries` was removed in its favor ([#2362](https://github.com/0xMiden/rust-sdk/pull/2362)).
@@ -53,6 +54,7 @@
 ### Enhancements
 
 * [FEATURE][rust] `Client::sync_state` now issues its independent gRPC calls concurrently instead of one after another, reducing the total time a sync takes. `NodeRpcClient::sync_notes_with_content` and `NodeRpcClient::sync_transactions` are now called concurrently rather than in sequence, and the per-account `NodeRpcClient::get_account` requests are issued in parallel instead of one at a time ([#2420](https://github.com/0xMiden/rust-sdk/pull/2420)).
+* [FEATURE][rust] Added `Client::track_account_witness`, `Client::untrack_account_witness` and `Client::tracked_account_witnesses` to register accounts whose account witness the sync keeps fresh in the store (new `account_witnesses` table). A transaction using a registered account as a foreign account builds its inputs from the store instead of issuing a `GetAccount` request, moving the cost from once per transaction to once per sync ([#2476](https://github.com/0xMiden/rust-sdk/pull/2476)).
 * [store] Added `SqliteStore::database_filepath`, which returns the backing database path losslessly as a `&Path` ([#2363](https://github.com/0xMiden/rust-sdk/pull/2363)).
 * [FEATURE][rust] Syncing no longer issues a `GetNotesById` request for a note whose attachments the `SyncNotes` response already carried. The node sends an attachment that fits in a single word verbatim, so the client reconstructs it locally: a private note whose attachments are all single-word is resolved with no follow-up request during state sync, as is a note of either type while checking expected notes. Both standard attachment schemes in `miden-standards` (`NetworkAccountTarget` and the PSWAP attachment) are single-word, so the round trip disappears from the common case. A caller-supplied attachment spanning more than one word arrives as a commitment and is still fetched. Public notes are unaffected during state sync, since their bodies are requested regardless ([#2360](https://github.com/0xMiden/rust-sdk/issues/2360)).
 * [FEATURE][rust] `CommittedNote` now carries the attachment content its source reported. `attachments()` returns it, `needs_attachment_fetch()` reports whether the content still has to be fetched via `GetNotesById`, and `with_attachments()` records it, rejecting content that does not hash to the metadata's attachments commitment ([#2360](https://github.com/0xMiden/rust-sdk/issues/2360)).

crates/rust-client/src/account/mod.rs

@@ -397,6 +397,41 @@ impl<AUTH> Client<AUTH> {
         self.add_account_inner(&account, ClientAccountType::Watched, true).await
     }
 
+    // ACCOUNT WITNESS PREFETCHING
+    // --------------------------------------------------------------------------------------------
+
+    /// Registers an account whose account witness [`Client::sync_chain`] keeps up to date, so that
+    /// transactions using it as a foreign account resolve the witness locally. This trades one
+    /// request per transaction for one per sync.
+    ///
+    /// A [`ForeignAccount::Private`](crate::transaction::ForeignAccount) needs nothing else, since
+    /// the caller supplies the account data. A
+    /// [`ForeignAccount::Public`](crate::transaction::ForeignAccount) additionally has to be
+    /// tracked by this client, so that its code, storage and vault come from the store as well;
+    /// registering an untracked public account costs a request per sync and saves none.
+    ///
+    /// The account is not validated against the network here. Registering an already registered
+    /// account is a no-op and keeps any cached witness.
+    pub async fn track_account_witness(&self, account_id: AccountId) -> Result<(), ClientError> {
+        self.store.track_account_witness(account_id).await.map_err(Into::into)
+    }
+
+    /// Stops keeping the account's witness up to date and drops the cached one.
+    ///
+    /// Returns `true` if the account was registered. Transactions using it keep working, falling
+    /// back to fetching the witness from the node.
+    pub async fn untrack_account_witness(
+        &self,
+        account_id: AccountId,
+    ) -> Result<bool, ClientError> {
+        self.store.untrack_account_witness(account_id).await.map_err(Into::into)
+    }
+
+    /// Returns the IDs of every account registered via [`Client::track_account_witness`].
+    pub async fn tracked_account_witnesses(&self) -> Result<Vec<AccountId>, ClientError> {
+        self.store.tracked_account_witnesses().await.map_err(Into::into)
+    }
+
     /// Fetches a public [`Account`] from the network, returning a typed error when the account
     /// doesn't exist on chain or is private.
     async fn fetch_public_account(&self, account_id: AccountId) -> Result<Account, ClientError> {

crates/rust-client/src/lib.rs

@@ -230,6 +230,7 @@ pub mod auth {
 
 /// Provides types for working with blocks within the Miden network.
 pub mod block {
+    pub use miden_protocol::block::account_tree::AccountWitness;
     pub use miden_protocol::block::{BlockHeader, BlockNumber, FeeParameters, ValidatorKeys};
 }
 

crates/rust-client/src/store/mod.rs

@@ -40,6 +40,7 @@ use miden_protocol::account::{
 };
 use miden_protocol::address::Address;
 use miden_protocol::asset::{Asset, AssetId, AssetVault, AssetWitness};
+use miden_protocol::block::account_tree::AccountWitness;
 use miden_protocol::block::{BlockHeader, BlockNumber};
 use miden_protocol::crypto::merkle::MerkleError;
 use miden_protocol::crypto::merkle::mmr::{Forest, InOrderIndex, MmrPeaks, PartialMmr};
@@ -471,6 +472,50 @@ pub trait Store: Send + Sync {
     /// Tag removal is the caller's responsibility — see [`Self::remove_note_tag`].
     async fn remove_address(&self, address: Address) -> Result<bool, StoreError>;
 
+    // ACCOUNT WITNESSES
+    // --------------------------------------------------------------------------------------------
+
+    /// Registers an account whose [`AccountWitness`] should be refreshed on every sync, so that
+    /// transactions using it as a foreign account can resolve the witness locally.
+    ///
+    /// No-op if the account is already registered; a cached witness is left in place. The witness
+    /// itself is filled in by the next sync.
+    async fn track_account_witness(&self, account_id: AccountId) -> Result<(), StoreError>;
+
+    /// Stops refreshing the account's witness and drops any cached one.
+    ///
+    /// Returns `true` if the account was registered.
+    async fn untrack_account_witness(&self, account_id: AccountId) -> Result<bool, StoreError>;
+
+    /// Retrieves the ID of every registered account, whether or not a witness has been cached for
+    /// it yet.
+    async fn tracked_account_witnesses(&self) -> Result<Vec<AccountId>, StoreError>;
+
+    /// Retrieves the cached [`AccountWitness`] along with the block it was fetched at.
+    ///
+    /// Callers must reject a witness whose block is not the one they execute against.
+    ///
+    /// Returns `None` when the account is not registered or has not been refreshed yet.
+    async fn get_account_witness(
+        &self,
+        account_id: AccountId,
+    ) -> Result<Option<(AccountWitness, BlockNumber)>, StoreError>;
+
+    /// Caches an [`AccountWitness`] for a registered account, replacing any previous one.
+    ///
+    /// Returns `false` if the account is not registered, in which case nothing is written.
+    /// Registering is [`Self::track_account_witness`]'s job alone.
+    ///
+    /// The caller should verify the witness against `block_num`'s account root first. The read
+    /// path only checks the block number, so a bad witness stored here surfaces later as a kernel
+    /// assertion during execution rather than as a chain validation error at sync time.
+    async fn update_account_witness(
+        &self,
+        account_id: AccountId,
+        witness: &AccountWitness,
+        block_num: BlockNumber,
+    ) -> Result<bool, StoreError>;
+
     // SETTINGS
     // --------------------------------------------------------------------------------------------
 

crates/rust-client/src/sync/mod.rs

@@ -61,16 +61,19 @@ use alloc::sync::Arc;
 use alloc::vec::Vec;
 use core::cmp::max;
 
+use futures::StreamExt;
 use miden_protocol::account::AccountId;
-use miden_protocol::block::BlockNumber;
+use miden_protocol::block::{BlockHeader, BlockNumber};
 use miden_protocol::crypto::merkle::mmr::{InOrderIndex, PartialMmr};
 use miden_protocol::note::NoteId;
 use miden_protocol::transaction::TransactionId;
 use miden_tx::auth::TransactionAuthenticator;
 use miden_tx::utils::serde::{Deserializable, DeserializationError, Serializable};
-use tracing::{debug, info};
+use tracing::{debug, info, warn};
 
 use crate::pswap::PswapChainObserver;
+use crate::rpc::AccountStateAt;
+use crate::rpc::domain::account::GetAccountRequest;
 use crate::store::{NoteFilter, TransactionFilter};
 use crate::{Client, ClientError};
 mod block_header;
@@ -82,6 +85,7 @@ mod note_observer;
 pub use note_observer::NoteObserver;
 
 mod state_sync;
+pub(crate) use state_sync::{MAX_CONCURRENT_ACCOUNT_FETCHES, validate_account_witness};
 pub use state_sync::{NoteUpdateAction, OnNoteReceived, StateSync, StateSyncInput};
 
 mod state_sync_update;
@@ -166,11 +170,92 @@ where
         // Cache MMR so pruning can reuse in-memory MMR.
         self.cache_partial_mmr(partial_mmr).await?;
 
+        self.refresh_account_witnesses().await;
+
         self.maybe_untrack_and_prune_irrelevant_blocks().await?;
 
         Ok(sync_summary)
     }
 
+    /// Refreshes the account witness of every registered account.
+    ///
+    /// Requests each witness at the new sync height, which is the reference block transactions
+    /// will execute against. Every account is refreshed whether or not its own state changed,
+    /// since a witness breaks when any other account in the tree moves.
+    ///
+    /// A failed refresh is logged and skipped rather than failing the sync: the stale entry stays,
+    /// and the read path rejects it by its block number.
+    async fn refresh_account_witnesses(&self) {
+        let account_ids = match self.store.tracked_account_witnesses().await {
+            Ok(account_ids) if account_ids.is_empty() => return,
+            Ok(account_ids) => account_ids,
+            Err(err) => {
+                warn!(%err, "failed to read the tracked account witness registry");
+                return;
+            },
+        };
+
+        let chain_tip_header = match self.get_latest_block_header().await {
+            Ok(header) => header,
+            Err(err) => {
+                warn!(%err, "failed to read the synced block header; skipping witness refresh");
+                return;
+            },
+        };
+
+        // Bounded fan-out, under the same limit the sync uses for its own `get_account` requests.
+        // Each future resolves to a `Result` so that one failure does not cancel the others.
+        let header = &chain_tip_header;
+        futures::stream::iter(account_ids)
+            .map(|account_id| async move {
+                (account_id, self.fetch_and_cache_account_witness(account_id, header).await)
+            })
+            .buffered(MAX_CONCURRENT_ACCOUNT_FETCHES)
+            .for_each(|(account_id, result)| async move {
+                if let Err(err) = result {
+                    warn!(%account_id, %err, "failed to refresh the cached account witness");
+                }
+            })
+            .await;
+    }
+
+    /// Fetches and stores a single account's witness at `chain_tip_header`'s block.
+    ///
+    /// Returns without a request when the applied sync update already left a witness at that
+    /// block, which it does for every public account it had to query for its state.
+    async fn fetch_and_cache_account_witness(
+        &self,
+        account_id: AccountId,
+        chain_tip_header: &BlockHeader,
+    ) -> Result<(), ClientError> {
+        let chain_tip = chain_tip_header.block_num();
+
+        if let Some((_, cached_at)) = self.store.get_account_witness(account_id).await?
+            && cached_at == chain_tip
+        {
+            return Ok(());
+        }
+
+        // The minimal request: no vault, no storage map entries, only the witness is wanted.
+        let (proof_block_num, proof) = self
+            .rpc_api
+            .get_account(account_id, GetAccountRequest::new().at(AccountStateAt::Block(chain_tip)))
+            .await?;
+
+        if proof_block_num != chain_tip {
+            return Err(ClientError::ChainValidationError(format!(
+                "get_account returned a proof at block {proof_block_num}, expected {chain_tip}"
+            )));
+        }
+
+        let (witness, _) = proof.into_parts();
+        validate_account_witness(&witness, account_id, chain_tip_header)?;
+
+        self.store.update_account_witness(account_id, &witness, chain_tip).await?;
+
+        Ok(())
+    }
+
     /// Fetches private notes from the Note Transport Layer for the tracked note tags.
     ///
     /// Returns the IDs of notes imported in this call. No-op (returns an empty vec) if note