diff --git a/crates/rust-client/src/note/mod.rs b/crates/rust-client/src/note/mod.rs index 1400f3e6d..ddbcb1ccd 100644 --- a/crates/rust-client/src/note/mod.rs +++ b/crates/rust-client/src/note/mod.rs @@ -139,7 +139,11 @@ where ) -> Result)>, 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, 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(¬e.clone().try_into()?) .await diff --git a/crates/rust-client/src/note/note_screener.rs b/crates/rust-client/src/note/note_screener.rs index 24e0d06b4..3241fac2e 100644 --- a/crates/rust-client/src/note/note_screener.rs +++ b/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,6 +57,8 @@ pub struct NoteScreener { store: Arc, /// A reference to the transaction authenticator authenticator: Option>, + /// Optional RPC client for lazy loading of data not found in local store. + rpc_client: Option>, } impl NoteScreener @@ -63,7 +66,19 @@ where AUTH: TransactionAuthenticator + Sync, { pub fn new(store: Arc, authenticator: Option>) -> Self { - Self { store, authenticator } + Self { store, authenticator, rpc_client: None } + } + + pub fn with_rpc( + store: Arc, + authenticator: Option>, + rpc_client: Arc, + ) -> 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()); diff --git a/crates/rust-client/src/store/data_store.rs b/crates/rust-client/src/store/data_store.rs index 20699ecbe..f5697f2a0 100644 --- a/crates/rust-client/src/store/data_store.rs +++ b/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, /// Cache of foreign account inputs that should be returned to the executor on demand. foreign_account_inputs: RwLock>, + /// Optional RPC client for lazy loading of data not found in local store. + rpc_client: Option>, } 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) -> 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, + rpc_client: Arc, + ) -> Self { + Self { + store, + transaction_mast_store: Arc::new(TransactionMastStore::new()), + foreign_account_inputs: RwLock::new(BTreeMap::new()), + rpc_client: Some(rpc_client), } } @@ -179,14 +202,84 @@ impl DataStore for ClientDataStore { foreign_account_id: AccountId, _ref_block: BlockNumber, ) -> Result { - 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( @@ -194,17 +287,44 @@ impl DataStore for ClientDataStore { script_root: Word, ) -> impl FutureMaybeSend, 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(¬e_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", + ))) + } + }, } } } diff --git a/crates/rust-client/src/sync/mod.rs b/crates/rust-client/src/sync/mod.rs index c9baa9086..7787680fd 100644 --- a/crates/rust-client/src/sync/mod.rs +++ b/crates/rust-client/src/sync/mod.rs @@ -122,7 +122,11 @@ where pub async fn sync_state(&mut self) -> Result { _ = 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); diff --git a/crates/rust-client/src/transaction/mod.rs b/crates/rust-client/src/transaction/mod.rs index c4228503d..204f1e447 100644 --- a/crates/rust-client/src/transaction/mod.rs +++ b/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, 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)?)