tracking for new StateSync component changes

bin/miden-cli/src/commands/sync.rs

@@ -12,8 +12,7 @@ impl SyncCmd {
         let new_details = client.sync_state().await?;
 
         println!("State synced to block {}", new_details.block_num);
-        println!("New public notes: {}", new_details.received_notes.len());
-        println!("Tracked notes updated: {}", new_details.committed_notes.len());
+        println!("Committed notes: {}", new_details.committed_notes.len());
         println!("Tracked notes consumed: {}", new_details.consumed_notes.len());
         println!("Tracked accounts updated: {}", new_details.updated_accounts.len());
         println!("Locked accounts: {}", new_details.locked_accounts.len());

bin/miden-cli/src/lib.rs

@@ -112,8 +112,8 @@ impl Cli {
         let authenticator = ClientAuthenticator::new(rng, keystore.clone());
 
         let client = Client::new(
-            Box::new(TonicRpcClient::new(
-                &(cli_config.rpc.endpoint.clone().into()),
+            Arc::new(TonicRpcClient::new(
+                &cli_config.rpc.endpoint.clone().into(),
                 cli_config.rpc.timeout_ms,
             )),
             rng,

bin/miden-cli/src/tests.rs

@@ -3,6 +3,7 @@ use std::{
     fs::File,
     io::{Read, Write},
     path::{Path, PathBuf},
+    sync::Arc,
 };
 
 use assert_cmd::Command;
@@ -738,7 +739,7 @@ async fn create_test_client_with_store_path(store_path: &Path) -> (TestClient, F
 
     let store = {
         let sqlite_store = SqliteStore::new(PathBuf::from(store_path)).await.unwrap();
-        std::sync::Arc::new(sqlite_store)
+        Arc::new(sqlite_store)
     };
 
     let mut rng = rand::thread_rng();
@@ -751,10 +752,10 @@ async fn create_test_client_with_store_path(store_path: &Path) -> (TestClient, F
     let authenticator = ClientAuthenticator::new(rng, keystore.clone());
     (
         TestClient::new(
-            Box::new(TonicRpcClient::new(&rpc_config.endpoint.into(), rpc_config.timeout_ms)),
+            Arc::new(TonicRpcClient::new(&rpc_config.endpoint.into(), rpc_config.timeout_ms)),
             rng,
             store,
-            std::sync::Arc::new(authenticator),
+            Arc::new(authenticator),
             true,
         ),
         keystore,

crates/rust-client/Cargo.toml

@@ -25,7 +25,7 @@ idxdb = ["dep:base64", "dep:serde-wasm-bindgen", "dep:wasm-bindgen", "dep:wasm-b
 sqlite = ["dep:rusqlite", "dep:deadpool-sqlite", "std"]
 std = ["miden-objects/std","miden-proving-service-client/std"]
 testing = ["miden-objects/testing", "miden-lib/testing", "miden-tx/testing"]
-tonic = ["std", "tonic/transport", "tonic/tls", "tonic/tls-native-roots"]
+tonic = ["std", "tonic/transport", "tonic/tls", "tonic/tls-native-roots", "dep:tokio"]
 web-tonic = ["dep:tonic-web-wasm-client"]
 
 [dependencies]

crates/rust-client/src/builder.rs

@@ -214,7 +214,7 @@ impl ClientBuilder {
         let authenticator = ClientAuthenticator::new(rng, keystore);
 
         Ok(Client::new(
-            rpc_api,
+            rpc_api.into(),
             rng,
             arc_store,
             Arc::new(authenticator),

crates/rust-client/src/lib.rs

@@ -83,7 +83,7 @@
 //! // Instantiate the client using a Tonic RPC client
 //! let endpoint = Endpoint::new("https".into(), "localhost".into(), Some(57291));
 //! let client: Client<RpoRandomCoin> = Client::new(
-//!     Box::new(TonicRpcClient::new(&endpoint, 10_000)),
+//!     Arc::new(TonicRpcClient::new(&endpoint, 10_000)),
 //!     rng,
 //!     store,
 //!     Arc::new(authenticator),
@@ -102,8 +102,6 @@
 #[macro_use]
 extern crate alloc;
 
-use alloc::boxed::Box;
-
 #[cfg(feature = "std")]
 extern crate std;
 
@@ -221,7 +219,7 @@ pub struct Client<R: FeltRng> {
     rng: R,
     /// An instance of [`NodeRpcClient`] which provides a way for the client to connect to the
     /// Miden node.
-    rpc_api: Box<dyn NodeRpcClient + Send>,
+    rpc_api: Arc<dyn NodeRpcClient + Send>,
     /// An instance of a [`LocalTransactionProver`] which will be the default prover for the
     /// client.
     tx_prover: Arc<LocalTransactionProver>,
@@ -257,7 +255,7 @@ impl<R: FeltRng> Client<R> {
     ///
     /// Returns an error if the client couldn't be instantiated.
     pub fn new(
-        rpc_api: Box<dyn NodeRpcClient + Send>,
+        rpc_api: Arc<dyn NodeRpcClient + Send>,
         rng: R,
         store: Arc<dyn Store>,
         authenticator: Arc<dyn TransactionAuthenticator>,
@@ -298,7 +296,7 @@ impl<R: FeltRng> Client<R> {
     // --------------------------------------------------------------------------------------------
 
     #[cfg(any(test, feature = "testing"))]
-    pub fn test_rpc_api(&mut self) -> &mut Box<dyn NodeRpcClient + Send> {
+    pub fn test_rpc_api(&mut self) -> &mut Arc<dyn NodeRpcClient + Send> {
         &mut self.rpc_api
     }
 

crates/rust-client/src/mock.rs

@@ -199,7 +199,7 @@ use alloc::boxed::Box;
 #[async_trait(?Send)]
 impl NodeRpcClient for MockRpcApi {
     async fn sync_notes(
-        &mut self,
+        &self,
         _block_num: BlockNumber,
         _note_tags: &[NoteTag],
     ) -> Result<NoteSyncInfo, RpcError> {
@@ -215,7 +215,7 @@ impl NodeRpcClient for MockRpcApi {
 
     /// Executes the specified sync state request and returns the response.
     async fn sync_state(
-        &mut self,
+        &self,
         block_num: BlockNumber,
         _account_ids: &[AccountId],
         _note_tags: &[NoteTag],
@@ -229,7 +229,7 @@ impl NodeRpcClient for MockRpcApi {
     /// Creates and executes a [GetBlockHeaderByNumberRequest].
     /// Only used for retrieving genesis block right now so that's the only case we need to cover.
     async fn get_block_header_by_number(
-        &mut self,
+        &self,
         block_num: Option<BlockNumber>,
         include_mmr_proof: bool,
     ) -> Result<(BlockHeader, Option<MmrProof>), RpcError> {
@@ -251,7 +251,7 @@ impl NodeRpcClient for MockRpcApi {
         Ok((block.header(), mmr_proof))
     }
 
-    async fn get_notes_by_id(&mut self, note_ids: &[NoteId]) -> Result<Vec<NetworkNote>, RpcError> {
+    async fn get_notes_by_id(&self, note_ids: &[NoteId]) -> Result<Vec<NetworkNote>, RpcError> {
         // assume all private notes for now
         let hit_notes = note_ids.iter().filter_map(|id| self.notes.get(id));
         let mut return_notes = vec![];
@@ -266,31 +266,28 @@ impl NodeRpcClient for MockRpcApi {
     }
 
     async fn submit_proven_transaction(
-        &mut self,
+        &self,
         _proven_transaction: ProvenTransaction,
     ) -> std::result::Result<(), RpcError> {
         // TODO: add some basic validations to test error cases
         Ok(())
     }
 
-    async fn get_account_update(
-        &mut self,
-        _account_id: AccountId,
-    ) -> Result<AccountDetails, RpcError> {
+    async fn get_account_update(&self, _account_id: AccountId) -> Result<AccountDetails, RpcError> {
         panic!("shouldn't be used for now")
     }
 
     async fn get_account_proofs(
-        &mut self,
-        _account_ids: &BTreeSet<ForeignAccount>,
+        &self,
+        _: &BTreeSet<ForeignAccount>,
         _code_commitments: Vec<AccountCode>,
     ) -> Result<AccountProofs, RpcError> {
         // TODO: Implement fully
         Ok((self.blocks.last().unwrap().header().block_num(), vec![]))
     }
 
     async fn check_nullifiers_by_prefix(
-        &mut self,
+        &self,
         _prefix: &[u16],
         _block_num: BlockNumber,
     ) -> Result<Vec<NullifierUpdate>, RpcError> {
@@ -315,9 +312,9 @@ pub async fn create_test_client() -> (MockClient, MockRpcApi, FilesystemKeyStore
 
     let authenticator = ClientAuthenticator::new(rng, keystore.clone());
     let rpc_api = MockRpcApi::new();
-    let boxed_rpc_api = Box::new(rpc_api.clone());
+    let arc_rpc_api = Arc::new(rpc_api.clone());
 
-    let client = MockClient::new(boxed_rpc_api, rng, store, Arc::new(authenticator.clone()), true);
+    let client = MockClient::new(arc_rpc_api, rng, store, Arc::new(authenticator.clone()), true);
     (client, rpc_api, keystore)
 }
 

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

@@ -94,7 +94,7 @@ impl<R: FeltRng> Client<R> {
     /// - If the note doesn't exist on the node.
     /// - If the note exists but is private.
     async fn import_note_record_by_id(
-        &mut self,
+        &self,
         previous_note: Option<InputNoteRecord>,
         id: NoteId,
     ) -> Result<Option<InputNoteRecord>, ClientError> {
@@ -136,7 +136,7 @@ impl<R: FeltRng> Client<R> {
     /// If the note isn't consumed and it was committed in the past relative to the client, then
     /// the MMR for the relevant block is fetched from the node and stored.
     async fn import_note_record_by_proof(
-        &mut self,
+        &self,
         previous_note: Option<InputNoteRecord>,
         note: Note,
         inclusion_proof: NoteInclusionProof,

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

@@ -56,7 +56,11 @@
 //! For more details on the API and error handling, see the documentation for the specific functions
 //! and types in this module.
 
-use alloc::{collections::BTreeSet, string::ToString, vec::Vec};
+use alloc::{
+    collections::{BTreeMap, BTreeSet},
+    string::ToString,
+    vec::Vec,
+};
 
 use miden_lib::transaction::TransactionKernel;
 use miden_objects::{account::AccountId, crypto::rand::FeltRng};
@@ -240,101 +244,88 @@ pub async fn get_input_note_with_id_prefix<R: FeltRng>(
 // ------------------------------------------------------------------------------------------------
 
 /// Contains note changes to apply to the store.
+#[derive(Clone, Debug, Default)]
 pub struct NoteUpdates {
-    /// A list of new input notes.
-    new_input_notes: Vec<InputNoteRecord>,
-    /// A list of new output notes.
-    new_output_notes: Vec<OutputNoteRecord>,
-    /// A list of updated input note records corresponding to locally-tracked input notes.
-    updated_input_notes: Vec<InputNoteRecord>,
-    /// A list of updated output note records corresponding to locally-tracked output notes.
-    updated_output_notes: Vec<OutputNoteRecord>,
+    /// 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(
-        new_input_notes: Vec<InputNoteRecord>,
-        new_output_notes: Vec<OutputNoteRecord>,
-        updated_input_notes: Vec<InputNoteRecord>,
-        updated_output_notes: Vec<OutputNoteRecord>,
+        updated_input_notes: impl IntoIterator<Item = InputNoteRecord>,
+        updated_output_notes: impl IntoIterator<Item = OutputNoteRecord>,
     ) -> Self {
         Self {
-            new_input_notes,
-            new_output_notes,
-            updated_input_notes,
-            updated_output_notes,
+            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(),
         }
     }
 
-    /// Combines two [`NoteUpdates`] into a single one.
-    #[must_use]
-    pub fn combine_with(mut self, other: Self) -> Self {
-        self.new_input_notes.extend(other.new_input_notes);
-        self.new_output_notes.extend(other.new_output_notes);
-        self.updated_input_notes.extend(other.updated_input_notes);
-        self.updated_output_notes.extend(other.updated_output_notes);
-
-        self
-    }
-
-    /// Returns all new input note records, meant to be tracked by the client.
-    pub fn new_input_notes(&self) -> &[InputNoteRecord] {
-        &self.new_input_notes
-    }
-
-    /// Returns all new output note records, meant to be tracked by the client.
-    pub fn new_output_notes(&self) -> &[OutputNoteRecord] {
-        &self.new_output_notes
-    }
-
-    /// Returns all updated input note records. That is, any input notes that are locally tracked
-    /// and have been updated.
-    pub fn updated_input_notes(&self) -> &[InputNoteRecord] {
-        &self.updated_input_notes
+    /// 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 updated output note records. That is, any output notes that are locally tracked
-    /// and have been updated.
-    pub fn updated_output_notes(&self) -> &[OutputNoteRecord] {
-        &self.updated_output_notes
+    /// 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()
-            && self.new_input_notes.is_empty()
-            && self.new_output_notes.is_empty()
+        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.
+    /// 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
-            .iter()
+            .values()
             .filter_map(|note_record| note_record.is_committed().then_some(note_record.id()));
 
         let committed_input_note_ids = self
             .updated_input_notes
-            .iter()
+            .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
+    /// 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
-            .iter()
+            .values()
             .filter_map(|note_record| note_record.is_consumed().then_some(note_record.id()));
 
         let consumed_input_note_ids = self
             .updated_input_notes
-            .iter()
+            .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<_>>()