feat: Add RPC-enabled DataStore for lazy loading

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

@@ -139,7 +139,11 @@ where
     ) -> Result<Vec<(InputNoteRecord, Vec<NoteConsumability>)>, ClientError> {
         let committed_notes = self.store.get_input_notes(NoteFilter::Committed).await?;
 
-        let note_screener = NoteScreener::new(self.store.clone(), self.authenticator.clone());
+        let note_screener = NoteScreener::with_rpc(
+            self.store.clone(),
+            self.authenticator.clone(),
+            self.rpc_api.clone(),
+        );
 
         let mut relevant_notes = Vec::new();
         for input_note in committed_notes {
@@ -169,7 +173,11 @@ where
         &self,
         note: InputNoteRecord,
     ) -> Result<Vec<NoteConsumability>, ClientError> {
-        let note_screener = NoteScreener::new(self.store.clone(), self.authenticator.clone());
+        let note_screener = NoteScreener::with_rpc(
+            self.store.clone(),
+            self.authenticator.clone(),
+            self.rpc_api.clone(),
+        );
         note_screener
             .check_relevance(&note.clone().try_into()?)
             .await

crates/rust-client/src/note/note_screener.rs

@@ -14,6 +14,7 @@ use miden_tx::{NoteCheckerError, NoteConsumptionChecker, TransactionExecutor};
 use thiserror::Error;
 
 use crate::ClientError;
+use crate::rpc::NodeRpcClient;
 use crate::rpc::domain::note::CommittedNote;
 use crate::store::data_store::ClientDataStore;
 use crate::store::{InputNoteRecord, NoteFilter, Store, StoreError};
@@ -56,14 +57,28 @@ pub struct NoteScreener<AUTH> {
     store: Arc<dyn Store>,
     /// A reference to the transaction authenticator
     authenticator: Option<Arc<AUTH>>,
+    /// Optional RPC client for lazy loading of data not found in local store.
+    rpc_client: Option<Arc<dyn NodeRpcClient>>,
 }
 
 impl<AUTH> NoteScreener<AUTH>
 where
     AUTH: TransactionAuthenticator + Sync,
 {
     pub fn new(store: Arc<dyn Store>, authenticator: Option<Arc<AUTH>>) -> Self {
-        Self { store, authenticator }
+        Self { store, authenticator, rpc_client: None }
+    }
+
+    pub fn with_rpc(
+        store: Arc<dyn Store>,
+        authenticator: Option<Arc<AUTH>>,
+        rpc_client: Arc<dyn NodeRpcClient>,
+    ) -> Self {
+        Self {
+            store,
+            authenticator,
+            rpc_client: Some(rpc_client),
+        }
     }
 
     /// Returns a vector of tuples describing the relevance of the provided note to the
@@ -126,7 +141,11 @@ where
 
         let tx_args = transaction_request.clone().into_transaction_args(tx_script);
 
-        let data_store = ClientDataStore::new(self.store.clone());
+        let data_store = if let Some(rpc) = &self.rpc_client {
+            ClientDataStore::with_rpc(self.store.clone(), rpc.clone())
+        } else {
+            ClientDataStore::new(self.store.clone())
+        };
         let mut transaction_executor = TransactionExecutor::new(&data_store);
         if let Some(authenticator) = &self.authenticator {
             transaction_executor = transaction_executor.with_authenticator(authenticator.as_ref());

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

@@ -13,8 +13,11 @@ use miden_objects::vm::FutureMaybeSend;
 use miden_objects::{MastForest, Word};
 use miden_tx::{DataStore, DataStoreError, MastForestStore, TransactionMastStore};
 
+use super::{PartialBlockchainFilter, Store};
+use crate::rpc::NodeRpcClient;
 use super::{AccountStorageFilter, PartialBlockchainFilter, Store};
 use crate::store::StoreError;
+use crate::transaction::ForeignAccount;
 use crate::utils::RwLock;
 
 // DATA STORE
@@ -28,14 +31,34 @@ pub struct ClientDataStore {
     transaction_mast_store: Arc<TransactionMastStore>,
     /// Cache of foreign account inputs that should be returned to the executor on demand.
     foreign_account_inputs: RwLock<BTreeMap<AccountId, AccountInputs>>,
+    /// Optional RPC client for lazy loading of data not found in local store.
+    rpc_client: Option<Arc<dyn NodeRpcClient>>,
 }
 
 impl ClientDataStore {
+    /// Creates a new `ClientDataStore` with an optional RPC client for lazy loading.
+    ///
+    /// If an RPC client is provided, the data store will attempt to fetch missing data
+    /// (such as note scripts and foreign account data) from the network when not found locally.
     pub fn new(store: alloc::sync::Arc<dyn Store>) -> Self {
         Self {
             store,
             transaction_mast_store: Arc::new(TransactionMastStore::new()),
             foreign_account_inputs: RwLock::new(BTreeMap::new()),
+            rpc_client: None,
+        }
+    }
+
+    /// Creates a new `ClientDataStore` with an RPC client for lazy loading.
+    pub fn with_rpc(
+        store: alloc::sync::Arc<dyn Store>,
+        rpc_client: Arc<dyn NodeRpcClient>,
+    ) -> Self {
+        Self {
+            store,
+            transaction_mast_store: Arc::new(TransactionMastStore::new()),
+            foreign_account_inputs: RwLock::new(BTreeMap::new()),
+            rpc_client: Some(rpc_client),
         }
     }
 
@@ -179,32 +202,129 @@ impl DataStore for ClientDataStore {
         foreign_account_id: AccountId,
         _ref_block: BlockNumber,
     ) -> Result<AccountInputs, DataStoreError> {
-        let cache = self.foreign_account_inputs.read();
+        // First, check the cache
+        {
+            let cache = self.foreign_account_inputs.read();
+            if let Some(inputs) = cache.get(&foreign_account_id) {
+                return Ok(inputs.clone());
+            }
+        }
 
-        let inputs = cache
-            .get(&foreign_account_id)
-            .cloned()
-            .ok_or(DataStoreError::AccountNotFound(foreign_account_id))?;
+        // If not in cache and RPC client is available, try fetching from the network
+        if let Some(rpc) = &self.rpc_client {
+            // Try to fetch as a public account with empty storage requirements
+            // This will work for public accounts, but won't work for private accounts
+            // (which require PartialAccount to be provided upfront)
+            if foreign_account_id.is_public() {
+                let foreign_account = ForeignAccount::Public(
+                    foreign_account_id,
+                    crate::rpc::domain::account::AccountStorageRequirements::default(),
+                );
+
+                let known_account_codes = self
+                    .store
+                    .get_foreign_account_code(vec![foreign_account_id])
+                    .await
+                    .map_err(|err| {
+                        DataStoreError::other(format!("Failed to get foreign account code: {err}"))
+                    })?;
+
+                match rpc
+                    .get_account_proofs(
+                        &[foreign_account].into_iter().collect(),
+                        known_account_codes,
+                    )
+                    .await
+                {
+                    Ok((_block_num, account_proofs)) => {
+                        if let Some(account_proof) = account_proofs
+                            .into_iter()
+                            .find(|proof| proof.account_id() == foreign_account_id)
+                        {
+                            let account_inputs: AccountInputs =
+                                account_proof.try_into().map_err(|err| {
+                                    DataStoreError::other(format!(
+                                        "Failed to convert account proof to AccountInputs: {err}"
+                                    ))
+                                })?;
+
+                            // Cache the fetched account inputs for future use
+                            {
+                                let mut cache = self.foreign_account_inputs.write();
+                                cache.insert(foreign_account_id, account_inputs.clone());
+                            }
+
+                            // Update the foreign account code cache
+                            if let Err(err) = self
+                                .store
+                                .upsert_foreign_account_code(
+                                    foreign_account_id,
+                                    account_inputs.code().clone(),
+                                )
+                                .await
+                            {
+                                // Log but don't fail - we still have the account inputs to return
+                                let _ = err;
+                            }
+
+                            return Ok(account_inputs);
+                        }
+                    },
+                    Err(rpc_err) => {
+                        return Err(DataStoreError::other(format!(
+                            "Failed to fetch foreign account {foreign_account_id} via RPC: {rpc_err}",
+                        )));
+                    },
+                }
+            }
+        }
 
-        Ok(inputs)
+        Err(DataStoreError::AccountNotFound(foreign_account_id))
     }
 
     fn get_note_script(
         &self,
         script_root: Word,
     ) -> impl FutureMaybeSend<Result<Option<NoteScript>, DataStoreError>> {
         let store = self.store.clone();
+        let rpc_client = self.rpc_client.clone();
 
         async move {
-            if let Ok(note_script) = store.get_note_script(script_root).await {
-                Ok(Some(note_script))
-            } else {
-                // If no matching note found, return an error
-                // TODO: refactor to make RPC call to `GetNoteScriptByRoot` in case notes are not
-                // found https://github.com/0xMiden/miden-client/issues/1410
-                Err(DataStoreError::other(
-                    format!("Note script with root {script_root} not found",),
-                ))
+            // First, try to get the note script from the local store
+            match store.get_note_script(script_root).await {
+                Ok(note_script) => Ok(note_script),
+                Err(_) => {
+                    // If not found locally and RPC client is available, try fetching from the
+                    // network
+                    if let Some(rpc) = rpc_client {
+                        match rpc.get_note_script_by_root(script_root).await {
+                            Ok(note_script) => {
+                                // Cache the fetched script in the local store for future use.
+                                // Since we know the script wasn't in the local store
+                                // (get_note_script failed),
+                                // upsert should effectively be an insert. If it fails (e.g., due to
+                                // database issues or concurrent
+                                // writes), we continue anyway since caching is just an
+                                // optimization - we still have the valid script to return.
+                                if let Err(_err) = store
+                                    .upsert_note_scripts(core::slice::from_ref(&note_script))
+                                    .await
+                                {
+                                    // In a no_std environment, we can't easily log, so we just
+                                    // continue
+                                }
+                                Ok(note_script)
+                            },
+                            Err(rpc_err) => Err(DataStoreError::other(format!(
+                                "Note script with root {script_root} not found via RPC: {rpc_err}",
+                            ))),
+                        }
+                    } else {
+                        Err(DataStoreError::other(format!(
+                            "Note script with root {script_root} not found in local store",
+                        )))
+                    }
+                },
             }
         }
     }

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

@@ -122,7 +122,11 @@ where
     pub async fn sync_state(&mut self) -> Result<SyncSummary, ClientError> {
         _ = self.ensure_genesis_in_place().await?;
 
-        let note_screener = NoteScreener::new(self.store.clone(), self.authenticator.clone());
+        let note_screener = NoteScreener::with_rpc(
+            self.store.clone(),
+            self.authenticator.clone(),
+            self.rpc_api.clone(),
+        );
         let state_sync =
             StateSync::new(self.rpc_api.clone(), Arc::new(note_screener), self.tx_graceful_blocks);
 

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

@@ -252,7 +252,7 @@ where
 
         let ignore_invalid_notes = transaction_request.ignore_invalid_input_notes();
 
-        let data_store = ClientDataStore::new(self.store.clone());
+        let data_store = ClientDataStore::with_rpc(self.store.clone(), self.rpc_api.clone());
         data_store.register_foreign_account_inputs(foreign_account_inputs.iter().cloned());
         for fpi_account in &foreign_account_inputs {
             data_store.mast_store().load_account_code(fpi_account.code());
@@ -442,7 +442,7 @@ where
 
         let account: Account = account_record.try_into()?;
 
-        let data_store = ClientDataStore::new(self.store.clone());
+        let data_store = ClientDataStore::with_rpc(self.store.clone(), self.rpc_api.clone());
 
         data_store.register_foreign_account_inputs(foreign_account_inputs.iter().cloned());
 
@@ -492,7 +492,11 @@ where
 
         // New relevant input notes
         let mut new_input_notes = vec![];
-        let note_screener = NoteScreener::new(self.store.clone(), self.authenticator.clone());
+        let note_screener = NoteScreener::with_rpc(
+            self.store.clone(),
+            self.authenticator.clone(),
+            self.rpc_api.clone(),
+        );
 
         for note in notes_from_output(executed_tx.output_notes()) {
             // TODO: check_relevance() should have the option to take multiple notes
@@ -722,7 +726,7 @@ where
         tx_args: TransactionArgs,
     ) -> Result<InputNotes<InputNote>, ClientError> {
         loop {
-            let data_store = ClientDataStore::new(self.store.clone());
+            let data_store = ClientDataStore::with_rpc(self.store.clone(), self.rpc_api.clone());
 
             data_store.mast_store().load_account_code(account.code());
             let execution = NoteConsumptionChecker::new(&self.build_executor(&data_store)?)