refactor: optimize StateSyncUpdate

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

@@ -370,6 +370,19 @@ impl NoteUpdateTracker {
         input_note_unspent_nullifiers.chain(output_note_unspent_nullifiers)
     }
 
+    /// Returns the block numbers containing input notes that remain unspent after this update.
+    pub(crate) fn unspent_input_note_block_numbers(
+        &self,
+    ) -> impl Iterator<Item = BlockNumber> + '_ {
+        self.input_notes
+            .values()
+            .chain(self.expected_input_notes.values())
+            .filter(|update| !update.inner().is_consumed())
+            .filter_map(|update| {
+                update.inner().inclusion_proof().map(|proof| proof.location().block_num())
+            })
+    }
+
     /// Appends nullifiers to the per-account ordered nullifier list.
     ///
     /// Nullifiers from the same account must be in execution order; ordering across different

crates/rust-client/src/pswap/discovery.rs

@@ -37,7 +37,7 @@ pub(crate) async fn discover_pswap_rounds(
     chain_note_updates: &[ObservedPswapNote],
 ) -> Result<Vec<PswapLineageRoundUpdate>, PswapLineageError> {
     let consumed_note_ids: BTreeSet<NoteId> =
-        state_sync_update.note_updates.consumed_note_ids().collect();
+        state_sync_update.note_updates().consumed_note_ids().collect();
 
     if consumed_note_ids.is_empty() && chain_note_updates.is_empty() {
         return Ok(Vec::new());
@@ -54,7 +54,7 @@ pub(crate) async fn discover_pswap_rounds(
 
     // Commit-block note roots for inserting reconstructed notes as `Committed`.
     let block_headers: BTreeMap<BlockNumber, BlockHeader> = state_sync_update
-        .partial_blockchain_updates
+        .partial_blockchain_updates()
         .block_headers()
         .map(|(header, _)| (header.block_num(), header.clone()))
         .collect();

crates/rust-client/src/pswap/observer.rs

@@ -86,7 +86,7 @@ impl NoteObserver for PswapChainObserver {
 
         // Nothing observed AND nothing consumed — correlator has no work.
         if chain_note_updates.is_empty()
-            && sync_update.note_updates.consumed_note_ids().next().is_none()
+            && sync_update.note_updates().consumed_note_ids().next().is_none()
         {
             return Ok(());
         }

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

@@ -84,6 +84,7 @@ mod state_sync;
 pub use state_sync::{NoteUpdateAction, OnNoteReceived, StateSync, StateSyncInput};
 
 mod state_sync_update;
+use state_sync_update::untrack_blocks;
 pub use state_sync_update::{
     AccountUpdates,
     PartialBlockchainUpdates,
@@ -289,10 +290,7 @@ where
             // Rebuild the PartialMmr and untrack each block to collect the authentication node
             // indices that are no longer needed by any remaining tracked leaf.
             let mut partial_mmr = self.get_current_partial_mmr().await?;
-            for &block_pos in &to_untrack {
-                nodes_to_remove
-                    .extend(partial_mmr.untrack(block_pos).into_iter().map(|(idx, _)| idx));
-            }
+            nodes_to_remove = untrack_blocks(&mut partial_mmr, to_untrack.iter().copied());
 
             blocks_to_untrack = to_untrack
                 .iter()