refactor: revert unnecessary changes

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

@@ -170,7 +170,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.store.build_current_partial_mmr(true).await?;
+                let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
 
                 let block_header = self
                     .get_and_store_authenticated_block(block_height, &mut current_partial_mmr)
@@ -214,7 +214,7 @@ impl<R: FeltRng> Client<R> {
 
         match committed_note_data {
             Some((metadata, inclusion_proof)) => {
-                let mut current_partial_mmr = self.store.build_current_partial_mmr(true).await?;
+                let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
                 let block_header = self
                     .get_and_store_authenticated_block(
                         inclusion_proof.location().block_num(),

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

@@ -4,7 +4,7 @@
 use alloc::{collections::BTreeSet, string::ToString, vec::Vec};
 
 use miden_lib::transaction::TransactionKernel;
-use miden_objects::{accounts::AccountId, crypto::rand::FeltRng, transaction::InputNote};
+use miden_objects::{accounts::AccountId, crypto::rand::FeltRng};
 
 use crate::{
     store::{InputNoteRecord, NoteFilter, OutputNoteRecord},
@@ -189,7 +189,9 @@ pub async fn get_input_note_with_id_prefix<R: FeltRng>(
 /// Contains note changes to apply to the store.
 pub struct NoteUpdates {
     /// A list of new input notes to be tracked.
-    new_input_notes: Vec<InputNote>,
+    new_input_notes: Vec<InputNoteRecord>,
+    /// A list of new output notes to be tracked.
+    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.
@@ -199,12 +201,14 @@ pub struct NoteUpdates {
 impl NoteUpdates {
     /// Creates a [NoteUpdates].
     pub fn new(
-        new_input_notes: Vec<InputNote>,
+        new_input_notes: Vec<InputNoteRecord>,
+        new_output_notes: Vec<OutputNoteRecord>,
         updated_input_notes: Vec<InputNoteRecord>,
         updated_output_notes: Vec<OutputNoteRecord>,
     ) -> Self {
         Self {
             new_input_notes,
+            new_output_notes,
             updated_input_notes,
             updated_output_notes,
         }
@@ -220,10 +224,15 @@ impl NoteUpdates {
     }
 
     /// Returns all new input note records, meant to be tracked by the client.
-    pub fn new_input_notes(&self) -> &[InputNote] {
+    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] {

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

@@ -10,7 +10,7 @@ use core::fmt::Debug;
 use async_trait::async_trait;
 use miden_objects::{
     accounts::{Account, AccountCode, AccountHeader, AccountId, AuthSecretKey},
-    crypto::merkle::{InOrderIndex, MmrPeaks, PartialMmr},
+    crypto::merkle::{InOrderIndex, MmrPeaks},
     notes::{NoteId, NoteTag, Nullifier},
     BlockHeader, Digest, Word,
 };
@@ -208,44 +208,6 @@ pub trait Store: Send + Sync {
         has_client_notes: bool,
     ) -> Result<(), StoreError>;
 
-    /// Builds the current store view of the chain's [PartialMmr]. Because we want to add all new
-    /// authentication nodes that could come from applying the MMR updates, we need to track all
-    /// known leaves thus far.
-    ///
-    /// As part of the syncing process, we add the current block number so we don't need to
-    /// track it here.
-    async fn build_current_partial_mmr(
-        &self,
-        include_current_block: bool,
-    ) -> Result<PartialMmr, StoreError> {
-        let current_block_num = self.get_sync_height().await?;
-
-        let tracked_nodes = self.get_chain_mmr_nodes(ChainMmrNodeFilter::All).await?;
-        let current_peaks = self.get_chain_mmr_peaks_by_block_num(current_block_num).await?;
-
-        let track_latest = if current_block_num != 0 {
-            match self.get_block_header_by_num(current_block_num - 1).await {
-                Ok((_, previous_block_had_notes)) => Ok(previous_block_had_notes),
-                Err(StoreError::BlockHeaderNotFound(_)) => Ok(false),
-                Err(err) => Err(err),
-            }?
-        } else {
-            false
-        };
-
-        let mut current_partial_mmr =
-            PartialMmr::from_parts(current_peaks, tracked_nodes, track_latest);
-
-        if include_current_block {
-            let (current_block, has_client_notes) =
-                self.get_block_header_by_num(current_block_num).await?;
-
-            current_partial_mmr.add(current_block.hash(), has_client_notes);
-        }
-
-        Ok(current_partial_mmr)
-    }
-
     // ACCOUNT
     // --------------------------------------------------------------------------------------------
 

crates/rust-client/src/store/sqlite_store/notes.rs

@@ -16,7 +16,6 @@ use rusqlite::{named_params, params, params_from_iter, types::Value, Connection,
 use super::SqliteStore;
 use crate::{
     notes::NoteUpdates,
-    rpc::generated::note,
     store::{
         note_record::OutputNoteState, InputNoteRecord, InputNoteState, NoteFilter,
         OutputNoteRecord, StoreError,
@@ -599,12 +598,18 @@ pub(crate) fn apply_note_updates_tx(
     tx: &Transaction,
     note_updates: &NoteUpdates,
 ) -> Result<(), StoreError> {
-    for input_note in note_updates.new_input_notes().iter() {
-        upsert_input_note_tx(tx, &input_note.clone().into())?;
+    for input_note in
+        note_updates.new_input_notes().iter().chain(note_updates.updated_input_notes())
+    {
+        upsert_input_note_tx(tx, input_note)?;
     }
 
-    for input_note_record in note_updates.updated_input_notes().iter() {
-        upsert_input_note_tx(tx, input_note_record)?;
+    for output_note in note_updates
+        .new_output_notes()
+        .iter()
+        .chain(note_updates.updated_output_notes())
+    {
+        upsert_output_note_tx(tx, output_note)?;
     }
 
     Ok(())

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

@@ -11,7 +11,7 @@ use tracing::warn;
 use super::NoteUpdates;
 use crate::{
     notes::NoteScreener,
-    store::{NoteFilter, StoreError},
+    store::{ChainMmrNodeFilter, NoteFilter, StoreError},
     Client, ClientError,
 };
 
@@ -20,7 +20,7 @@ impl<R: FeltRng> Client<R> {
     /// Updates committed notes with no MMR data. These could be notes that were
     /// imported with an inclusion proof, but its block header isn't tracked.
     pub(crate) async fn update_mmr_data(&mut self) -> Result<(), ClientError> {
-        let mut current_partial_mmr = self.store.build_current_partial_mmr(true).await?;
+        let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
 
         let mut changed_notes = vec![];
         for mut note in self.store.get_input_notes(NoteFilter::Unverified).await? {
@@ -101,6 +101,44 @@ impl<R: FeltRng> Client<R> {
         Ok(false)
     }
 
+    /// Builds the current store view of the chain's [PartialMmr]. Because we want to add all new
+    /// authentication nodes that could come from applying the MMR updates, we need to track all
+    /// known leaves thus far.
+    ///
+    /// As part of the syncing process, we add the current block number so we don't need to
+    /// track it here.
+    pub(crate) async fn build_current_partial_mmr(
+        &self,
+        include_current_block: bool,
+    ) -> Result<PartialMmr, ClientError> {
+        let current_block_num = self.store.get_sync_height().await?;
+
+        let tracked_nodes = self.store.get_chain_mmr_nodes(ChainMmrNodeFilter::All).await?;
+        let current_peaks = self.store.get_chain_mmr_peaks_by_block_num(current_block_num).await?;
+
+        let track_latest = if current_block_num != 0 {
+            match self.store.get_block_header_by_num(current_block_num - 1).await {
+                Ok((_, previous_block_had_notes)) => Ok(previous_block_had_notes),
+                Err(StoreError::BlockHeaderNotFound(_)) => Ok(false),
+                Err(err) => Err(ClientError::StoreError(err)),
+            }?
+        } else {
+            false
+        };
+
+        let mut current_partial_mmr =
+            PartialMmr::from_parts(current_peaks, tracked_nodes, track_latest);
+
+        if include_current_block {
+            let (current_block, has_client_notes) =
+                self.store.get_block_header_by_num(current_block_num).await?;
+
+            current_partial_mmr.add(current_block.hash(), has_client_notes);
+        }
+
+        Ok(current_partial_mmr)
+    }
+
     /// Retrieves and stores a [BlockHeader] by number, and stores its authentication data as well.
     ///
     /// If the store already contains MMR data for the requested block number, the request isn't
@@ -158,7 +196,7 @@ impl<R: FeltRng> Client<R> {
             return self.ensure_genesis_in_place().await;
         }
 
-        let mut current_partial_mmr = self.store.build_current_partial_mmr(true).await?;
+        let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
         let anchor_block = self
             .get_and_store_authenticated_block(epoch_block_number, &mut current_partial_mmr)
             .await?;

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

@@ -160,7 +160,7 @@ impl<R: FeltRng> Client<R> {
                 note_tags,
                 unspent_input_notes,
                 unspent_output_notes,
-                self.store.build_current_partial_mmr(false).await?,
+                self.build_current_partial_mmr(false).await?,
             )
             .sync_state_step()
             .await?;

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

@@ -252,7 +252,7 @@ impl StateSync {
     /// * If the account is public, the component will request the node for the updated account
     ///   details.
     /// * If the account is private it will be marked as mismatched and the client will need to
-    ///  handle it (it could be a stale account state or a reason to lock the account).
+    ///   handle it (it could be a stale account state or a reason to lock the account).
     async fn account_state_sync(
         &self,
         account_hash_updates: &[(AccountId, Digest)],
@@ -288,10 +288,9 @@ impl StateSync {
     ///
     /// The transaction changes might include:
     /// * Transactions that were committed in the block. Some of these might me tracked by the
-    ///   client
-    ///  and need to be marked as committed.
+    ///   client and need to be marked as committed.
     /// * Local tracked transactions that were discarded because the notes that they were processing
-    /// were nullified by an another transaction.
+    ///   were nullified by an another transaction.
     async fn note_state_sync(
         &mut self,
         block_header: &BlockHeader,
@@ -316,7 +315,8 @@ impl StateSync {
             .filter_map(|note_id| self.unspent_output_notes.remove(note_id))
             .collect();
 
-        let note_updates = NoteUpdates::new(new_notes, modified_input_notes, modified_output_notes);
+        let note_updates =
+            NoteUpdates::new(new_notes, vec![], modified_input_notes, modified_output_notes);
         let transaction_updates =
             TransactionUpdates::new(committed_transactions, discarded_transactions);
 
@@ -329,7 +329,7 @@ impl StateSync {
         &mut self,
         committed_notes: Vec<CommittedNote>,
         block_header: &BlockHeader,
-    ) -> Result<Vec<InputNote>, ClientError> {
+    ) -> Result<Vec<InputNoteRecord>, ClientError> {
         let mut new_public_notes = vec![];
 
         // We'll only pick committed notes that we are tracking as input/output notes. Since the
@@ -369,8 +369,12 @@ impl StateSync {
         }
 
         // Query the node for input note data and build the entities
-        let new_public_notes =
-            self.fetch_public_note_details(&new_public_notes, block_header).await?;
+        let new_public_notes = self
+            .fetch_public_note_details(&new_public_notes, block_header)
+            .await?
+            .into_iter()
+            .map(|note| note.into())
+            .collect();
 
         Ok(new_public_notes)
     }

crates/rust-client/src/tests.rs

@@ -369,7 +369,7 @@ async fn test_sync_state_mmr() {
     );
 
     // Try reconstructing the chain_mmr from what's in the database
-    let partial_mmr = client.test_store().build_current_partial_mmr(true).await.unwrap();
+    let partial_mmr = client.build_current_partial_mmr(true).await.unwrap();
     assert_eq!(partial_mmr.forest(), 6);
     assert!(partial_mmr.open(0).unwrap().is_some()); // Account anchor block
     assert!(partial_mmr.open(1).unwrap().is_some());

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

@@ -832,7 +832,7 @@ impl<R: FeltRng> Client<R> {
             let summary = self.sync_state().await?;
 
             if summary.block_num != block_num {
-                let mut current_partial_mmr = self.store.build_current_partial_mmr(true).await?;
+                let mut current_partial_mmr = self.build_current_partial_mmr(true).await?;
                 self.get_and_store_authenticated_block(block_num, &mut current_partial_mmr)
                     .await?;
             }