feat: add StateSync component

CHANGELOG.md

@@ -48,6 +48,7 @@
 
 ### Changes
 
+* [BREAKING] Refactored the sync process to use a new `SyncState` component (#650).
 * [BREAKING] Return `None` instead of `Err` when an entity is not found (#632).
 * Add support for notes without assets in transaction requests (#654).
 * Refactored RPC functions and structs to improve code quality (#616).

bin/miden-cli/src/lib.rs

@@ -9,6 +9,7 @@ use miden_client::{
     crypto::{FeltRng, RpoRandomCoin},
     rpc::TonicRpcClient,
     store::{sqlite_store::SqliteStore, NoteFilter as ClientNoteFilter, OutputNoteRecord, Store},
+    sync::{on_note_received, StateSync},
     Client, ClientError, Felt, IdPrefixFetchError,
 };
 use rand::Rng;
@@ -111,14 +112,28 @@ impl Cli {
             .map_err(CliError::KeyStore)?;
         let authenticator = ClientAuthenticator::new(rng, keystore.clone());
 
+        let rpc_api = Arc::new(TonicRpcClient::new(
+            &cli_config.rpc.endpoint.clone().into(),
+            cli_config.rpc.timeout_ms,
+        ));
+
+        let state_sync_component = StateSync::new(
+            rpc_api.clone(),
+            Box::new({
+                let store_clone = store.clone();
+                move |committed_note, public_note| {
+                    Box::pin(on_note_received(store_clone.clone(), committed_note, public_note))
+                }
+            }),
+            Box::new(move |_committed_note| Box::pin(async { Ok(true) })),
+        );
+
         let client = Client::new(
-            Arc::new(TonicRpcClient::new(
-                &cli_config.rpc.endpoint.clone().into(),
-                cli_config.rpc.timeout_ms,
-            )),
+            rpc_api,
             rng,
             store as Arc<dyn Store>,
             Arc::new(authenticator),
+            state_sync_component,
             in_debug_mode,
         );
 

bin/miden-cli/src/tests.rs

@@ -22,6 +22,7 @@ use miden_client::{
     },
     rpc::{Endpoint, TonicRpcClient},
     store::{sqlite_store::SqliteStore, NoteFilter},
+    sync::{on_note_received, StateSync},
     testing::account_id::ACCOUNT_ID_OFF_CHAIN_SENDER,
     transaction::{OutputNote, TransactionRequestBuilder},
     utils::Serializable,
@@ -748,16 +749,23 @@ async fn create_test_client_with_store_path(store_path: &Path) -> (TestClient, F
     let rng = RpoRandomCoin::new(coin_seed.map(Felt::new));
 
     let keystore = FilesystemKeyStore::new(temp_dir()).unwrap();
-
     let authenticator = ClientAuthenticator::new(rng, keystore.clone());
+
+    let rpc_api = Arc::new(TonicRpcClient::new(&rpc_config.endpoint.into(), rpc_config.timeout_ms));
+
+    let state_sync_component = StateSync::new(
+        rpc_api.clone(),
+        Box::new({
+            let store_clone = store.clone();
+            move |committed_note, public_note| {
+                Box::pin(on_note_received(store_clone.clone(), committed_note, public_note))
+            }
+        }),
+        Box::new(move |_committed_note| Box::pin(async { Ok(true) })),
+    );
+
     (
-        TestClient::new(
-            Arc::new(TonicRpcClient::new(&rpc_config.endpoint.into(), rpc_config.timeout_ms)),
-            rng,
-            store,
-            Arc::new(authenticator),
-            true,
-        ),
+        TestClient::new(rpc_api, rng, store, Arc::new(authenticator), state_sync_component, true),
         keystore,
     )
 }

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

@@ -41,7 +41,7 @@ use miden_lib::account::{auth::RpoFalcon512, wallets::BasicWallet};
 use miden_objects::{
     block::BlockHeader,
     crypto::{dsa::rpo_falcon512::PublicKey, rand::FeltRng},
-    AccountError, Word,
+    AccountError, Digest, Word,
 };
 
 use super::Client;
@@ -300,6 +300,51 @@ pub fn build_wallet_id(
     Ok(account.id())
 }
 
+// ACCOUNT UPDATES
+// ================================================================================================
+
+#[derive(Debug, Clone, Default)]
+/// Contains account changes to apply to the store after a sync request.
+pub struct AccountUpdates {
+    /// Updated public accounts.
+    updated_public_accounts: Vec<Account>,
+    /// Account hashes received from the network that don't match the currently locally-tracked
+    /// state of the private accounts.
+    ///
+    /// These updates may represent a stale account hash (meaning that the latest local state
+    /// hasn't been committed). If this is not the case, the account may be locked until the state
+    /// is restored manually.
+    mismatched_private_accounts: Vec<(AccountId, Digest)>,
+}
+
+impl AccountUpdates {
+    /// Creates a new instance of `AccountUpdates`.
+    pub fn new(
+        updated_public_accounts: Vec<Account>,
+        mismatched_private_accounts: Vec<(AccountId, Digest)>,
+    ) -> Self {
+        Self {
+            updated_public_accounts,
+            mismatched_private_accounts,
+        }
+    }
+
+    /// Returns the updated public accounts.
+    pub fn updated_public_accounts(&self) -> &[Account] {
+        &self.updated_public_accounts
+    }
+
+    /// Returns the mismatched private accounts.
+    pub fn mismatched_private_accounts(&self) -> &[(AccountId, Digest)] {
+        &self.mismatched_private_accounts
+    }
+
+    pub fn extend(&mut self, other: AccountUpdates) {
+        self.updated_public_accounts.extend(other.updated_public_accounts);
+        self.mismatched_private_accounts.extend(other.mismatched_private_accounts);
+    }
+}
+
 // TESTS
 // ================================================================================================
 

crates/rust-client/src/lib.rs

@@ -192,6 +192,7 @@ use miden_tx::{
 };
 use rpc::NodeRpcClient;
 use store::{data_store::ClientDataStore, Store};
+use sync::StateSync;
 use tracing::info;
 
 // MIDEN CLIENT
@@ -218,6 +219,9 @@ pub struct Client<R: FeltRng> {
     tx_prover: Arc<LocalTransactionProver>,
     /// An instance of a [`TransactionExecutor`] that will be used to execute transactions.
     tx_executor: TransactionExecutor,
+    /// An instance of a [`StateSync`] component that will be used to synchronize the client's
+    /// state with the state of the Miden network.
+    state_sync_component: StateSync,
     /// Flag to enable the debug mode for scripts compilation and execution.
     in_debug_mode: bool,
 }
@@ -240,6 +244,8 @@ impl<R: FeltRng> Client<R> {
     ///   store as the one for `store`, but it doesn't have to be the **same instance**.
     /// - `authenticator`: Defines the transaction authenticator that will be used by the
     ///   transaction executor whenever a signature is requested from within the VM.
+    /// - `state_sync_component`: An instance of [`StateSync`] that will be used to synchronize the
+    ///   client's state with the state of the Miden network.
     /// - `in_debug_mode`: Instantiates the transaction executor (and in turn, its compiler) in
     ///   debug mode, which will enable debug logs for scripts compiled with this mode for easier
     ///   MASM debugging.
@@ -252,6 +258,7 @@ impl<R: FeltRng> Client<R> {
         rng: R,
         store: Arc<dyn Store>,
         authenticator: Arc<dyn TransactionAuthenticator>,
+        state_sync_component: StateSync,
         in_debug_mode: bool,
     ) -> Self {
         let data_store = Arc::new(ClientDataStore::new(store.clone())) as Arc<dyn DataStore>;
@@ -270,6 +277,7 @@ impl<R: FeltRng> Client<R> {
             rpc_api,
             tx_prover,
             tx_executor,
+            state_sync_component,
             in_debug_mode,
         }
     }

crates/rust-client/src/mock.rs

@@ -44,6 +44,7 @@ use crate::{
         NodeRpcClient, RpcError,
     },
     store::sqlite_store::SqliteStore,
+    sync::{on_note_received, StateSync},
     transaction::ForeignAccount,
     Client,
 };
@@ -320,11 +321,23 @@ pub async fn create_test_client() -> (MockClient, MockRpcApi, FilesystemKeyStore
 
     let keystore = FilesystemKeyStore::new(temp_dir()).unwrap();
 
-    let authenticator = ClientAuthenticator::new(rng, keystore.clone());
+    let authenticator = Arc::new(ClientAuthenticator::new(rng, keystore.clone()));
     let rpc_api = MockRpcApi::new();
     let arc_rpc_api = Arc::new(rpc_api.clone());
 
-    let client = MockClient::new(arc_rpc_api, rng, store, Arc::new(authenticator.clone()), true);
+    let state_sync_component = StateSync::new(
+        arc_rpc_api.clone(),
+        Box::new({
+            let store_clone = store.clone();
+            move |committed_note, public_note| {
+                Box::pin(on_note_received(store_clone.clone(), committed_note, public_note))
+            }
+        }),
+        Box::new(move |_committed_note| Box::pin(async { Ok(true) })),
+    );
+
+    let client =
+        MockClient::new(arc_rpc_api, rng, store, authenticator, state_sync_component, true);
     (client, rpc_api, keystore)
 }
 

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

@@ -169,7 +169,7 @@ impl<R: FeltRng> Client<R> {
                 note_record.inclusion_proof_received(inclusion_proof, metadata)?;
 
             if block_height < current_block_num {
-                let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
+                let mut current_partial_mmr = self.build_current_partial_mmr().await?;
 
                 let block_header = self
                     .get_and_store_authenticated_block(block_height, &mut current_partial_mmr)
@@ -213,7 +213,7 @@ impl<R: FeltRng> Client<R> {
 
         match committed_note_data {
             Some((metadata, inclusion_proof)) => {
-                let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
+                let mut current_partial_mmr = self.build_current_partial_mmr().await?;
                 let block_header = self
                     .get_and_store_authenticated_block(
                         inclusion_proof.location().block_num(),

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

@@ -56,11 +56,7 @@
 //! For more details on the API and error handling, see the documentation for the specific functions
 //! and types in this module.
 
-use alloc::{
-    collections::{BTreeMap, BTreeSet},
-    string::ToString,
-    vec::Vec,
-};
+use alloc::{string::ToString, vec::Vec};
 
 use miden_lib::transaction::TransactionKernel;
 use miden_objects::{account::AccountId, crypto::rand::FeltRng};
@@ -74,6 +70,7 @@ pub mod script_roots;
 
 mod import;
 mod note_screener;
+mod note_updates;
 
 // RE-EXPORTS
 // ================================================================================================
@@ -92,6 +89,7 @@ pub use miden_objects::{
     NoteError,
 };
 pub use note_screener::{NoteConsumability, NoteRelevance, NoteScreener, NoteScreenerError};
+pub use note_updates::NoteUpdates;
 
 /// Contains functions to simplify standard note scripts creation.
 pub mod scripts {
@@ -239,103 +237,3 @@ pub async fn get_input_note_with_id_prefix<R: FeltRng>(
         .pop()
         .expect("input_note_records should always have one element"))
 }
-
-// NOTE UPDATES
-// ------------------------------------------------------------------------------------------------
-
-/// Contains note changes to apply to the store.
-#[derive(Clone, Debug, Default)]
-pub struct NoteUpdates {
-    /// A map of new and updated input note records to be upserted in the store.
-    updated_input_notes: BTreeMap<NoteId, InputNoteRecord>,
-    /// A map of updated output note records to be upserted in the store.
-    updated_output_notes: BTreeMap<NoteId, OutputNoteRecord>,
-}
-
-impl NoteUpdates {
-    /// Creates a [`NoteUpdates`].
-    pub fn new(
-        updated_input_notes: impl IntoIterator<Item = InputNoteRecord>,
-        updated_output_notes: impl IntoIterator<Item = OutputNoteRecord>,
-    ) -> Self {
-        Self {
-            updated_input_notes: updated_input_notes
-                .into_iter()
-                .map(|note| (note.id(), note))
-                .collect(),
-            updated_output_notes: updated_output_notes
-                .into_iter()
-                .map(|note| (note.id(), note))
-                .collect(),
-        }
-    }
-
-    /// Returns all input note records that have been updated.
-    /// This may include:
-    /// - New notes that have been created that should be inserted.
-    /// - Existing tracked notes that should be updated.
-    pub fn updated_input_notes(&self) -> impl Iterator<Item = &InputNoteRecord> {
-        self.updated_input_notes.values()
-    }
-
-    /// Returns all output note records that have been updated.
-    /// This may include:
-    /// - New notes that have been created that should be inserted.
-    /// - Existing tracked notes that should be updated.
-    pub fn updated_output_notes(&self) -> impl Iterator<Item = &OutputNoteRecord> {
-        self.updated_output_notes.values()
-    }
-
-    /// Returns whether no new note-related information has been retrieved.
-    pub fn is_empty(&self) -> bool {
-        self.updated_input_notes.is_empty() && self.updated_output_notes.is_empty()
-    }
-
-    /// Returns any note that has been committed into the chain in this update (either new or
-    /// already locally tracked)
-    pub fn committed_input_notes(&self) -> impl Iterator<Item = &InputNoteRecord> {
-        self.updated_input_notes.values().filter(|note| note.is_committed())
-    }
-
-    /// Returns the IDs of all notes that have been committed in this update.
-    /// This includes both new notes and tracked expected notes that were committed in this update.
-    pub fn committed_note_ids(&self) -> BTreeSet<NoteId> {
-        let committed_output_note_ids = self
-            .updated_output_notes
-            .values()
-            .filter_map(|note_record| note_record.is_committed().then_some(note_record.id()));
-
-        let committed_input_note_ids = self
-            .updated_input_notes
-            .values()
-            .filter_map(|note_record| note_record.is_committed().then_some(note_record.id()));
-
-        committed_input_note_ids
-            .chain(committed_output_note_ids)
-            .collect::<BTreeSet<_>>()
-    }
-
-    /// Returns the IDs of all notes that have been consumed.
-    /// This includes both notes that have been consumed locally or externally in this update.
-    pub fn consumed_note_ids(&self) -> BTreeSet<NoteId> {
-        let consumed_output_note_ids = self
-            .updated_output_notes
-            .values()
-            .filter_map(|note_record| note_record.is_consumed().then_some(note_record.id()));
-
-        let consumed_input_note_ids = self
-            .updated_input_notes
-            .values()
-            .filter_map(|note_record| note_record.is_consumed().then_some(note_record.id()));
-
-        consumed_input_note_ids.chain(consumed_output_note_ids).collect::<BTreeSet<_>>()
-    }
-
-    /// Extends this note update information with `other`. If the two contain updates to the same
-    /// note (i.e. their IDs match), the updates from `other` will overwrite the updates in
-    /// `self`.
-    pub(crate) fn extend(&mut self, other: NoteUpdates) {
-        self.updated_input_notes.extend(other.updated_input_notes);
-        self.updated_output_notes.extend(other.updated_output_notes);
-    }
-}

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

@@ -57,6 +57,8 @@ impl NoteScreener {
     /// Does a fast check for known scripts (P2ID, P2IDR, SWAP). We're currently
     /// unable to execute notes that aren't committed so a slow check for other scripts is
     /// currently not available.
+    ///
+    /// If relevance can't be determined, the screener defaults to setting the note as consumable.
     pub async fn check_relevance(
         &self,
         note: &Note,