From 27905d94d6d29b1cc5468135fbb6580b832a26fa Mon Sep 17 00:00:00 2001
From: jonesmarvin8 <83104039+jonesmarvin8@users.noreply.github.com>
Date: Tue, 28 Apr 2026 14:13:04 -0400
Subject: [PATCH] addressed more comments

---
 wallet/src/chain_storage.rs                   |  6 +-
 wallet/src/lib.rs                             | 23 ++++--
 .../native_token_transfer/public.rs           | 70 ++++++++-----------
 3 files changed, 47 insertions(+), 52 deletions(-)

diff --git a/wallet/src/chain_storage.rs b/wallet/src/chain_storage.rs
index cf6de5ce..a3634367 100644
--- a/wallet/src/chain_storage.rs
+++ b/wallet/src/chain_storage.rs
@@ -56,9 +56,9 @@ impl WalletChainStore {
             .expect("Malformed persistent account data, must have private root");
 
         let mut public_tree = KeyTreePublic::new_from_root(match public_root {
-            PersistentAccountData::Public(data) => {
-                data.data.expect("Expect valid public account keys")
-            }
+            PersistentAccountData::Public(data) => data
+                .data
+                .expect("public tree in persistent_accounts failed to return a valid KeyTree."),
             _ => unreachable!(),
         });
         let mut private_tree = KeyTreePrivate::new_from_root(match private_root {
diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs
index 34633ced..fdb83d94 100644
--- a/wallet/src/lib.rs
+++ b/wallet/src/lib.rs
@@ -548,14 +548,14 @@ impl WalletCore {
     }
 
     pub fn sign_public_message(
-        wallet: &Self,
+        &self,
         message: &nssa::public_transaction::Message,
         account_ids: &[AccountId],
     ) -> Result<nssa::public_transaction::WitnessSet, ExecutionFailureKind> {
         let mut private_keys = Vec::new();
 
         for &account_id in account_ids {
-            let key = wallet
+            let key = self
                 .storage
                 .user_data
                 .get_pub_account_signing_key(account_id)
@@ -569,16 +569,25 @@ impl WalletCore {
         ))
     }
 
+    #[must_use]
     pub fn sign_privacy_message(
         message: &nssa::privacy_preserving_transaction::Message,
         proof: &Proof,
         acc_manager: &privacy_preserving_tx::AccountManager,
-    ) -> nssa::privacy_preserving_transaction::witness_set::WitnessSet
-    {
+    ) -> nssa::privacy_preserving_transaction::witness_set::WitnessSet {
         nssa::privacy_preserving_transaction::witness_set::WitnessSet::for_message(
-                message,
-                proof.clone(),
-                &acc_manager.public_account_auth())
+            message,
+            proof.clone(),
+            &acc_manager.public_account_auth(),
+        )
+    }
 
+    #[must_use]
+    pub fn filter_owned_accounts(&self, account_ids: &[nssa::AccountId]) -> Vec<nssa::AccountId> {
+        account_ids
+            .iter()
+            .filter(|&&account_id| self.get_account_public_signing_key(account_id).is_some())
+            .copied()
+            .collect()
     }
 }
diff --git a/wallet/src/program_facades/native_token_transfer/public.rs b/wallet/src/program_facades/native_token_transfer/public.rs
index 7705c268..2ddb0d87 100644
--- a/wallet/src/program_facades/native_token_transfer/public.rs
+++ b/wallet/src/program_facades/native_token_transfer/public.rs
@@ -22,49 +22,35 @@ impl NativeTokenTransfer<'_> {
             .await
             .map_err(ExecutionFailureKind::SequencerError)?;
 
-        if balance >= balance_to_move {
-            let account_ids = vec![from, to];
-            let program_id = Program::authenticated_transfer_program().id();
-
-            let mut sign_ids = Vec::new();
-            sign_ids.push(from);
-
-            let mut nonces = self
-                .0
-                .get_accounts_nonces(vec![from])
-                .await
-                .map_err(ExecutionFailureKind::SequencerError)?;
-            let to_signing_key = self.0.storage.user_data.get_pub_account_signing_key(to);
-            if let Some(_to_signing_key) = to_signing_key {
-                sign_ids.push(to);
-                let to_nonces = self
-                    .0
-                    .get_accounts_nonces(vec![to])
-                    .await
-                    .map_err(ExecutionFailureKind::SequencerError)?;
-                nonces.extend(to_nonces);
-            } else {
-                println!(
-                    "Receiver's account ({to}) private key not found in wallet. Proceeding with only sender's key."
-                );
-            }
-
-            let message =
-                Message::try_new(program_id, account_ids, nonces, balance_to_move).unwrap();
-
-            let witness_set = WalletCore::sign_public_message(self.0, &message, &sign_ids)
-                .expect("Expect a valid signature");
-
-            let tx = PublicTransaction::new(message, witness_set);
-
-            Ok(self
-                .0
-                .sequencer_client
-                .send_transaction(NSSATransaction::Public(tx))
-                .await?)
-        } else {
-            Err(ExecutionFailureKind::InsufficientFundsError)
+        if balance < balance_to_move {
+            return Err(ExecutionFailureKind::InsufficientFundsError);
         }
+
+        let account_ids = vec![from, to];
+        let program_id = Program::authenticated_transfer_program().id();
+
+        let sign_ids = self.0.filter_owned_accounts(&[from, to]);
+
+        // Fetch nonces for both accounts unconditionally
+        let nonces = self
+            .0
+            .get_accounts_nonces(account_ids.clone())
+            .await
+            .map_err(ExecutionFailureKind::SequencerError)?;
+
+        let message = Message::try_new(program_id, account_ids, nonces, balance_to_move).unwrap();
+
+        // Assumes this now silently skips accounts without signing keys
+        let witness_set = WalletCore::sign_public_message(self.0, &message, &sign_ids)
+            .expect("`WalletCore::sign_public_message() failed to produce a signature for a NativeTokenTransfer.");
+
+        let tx = PublicTransaction::new(message, witness_set);
+
+        Ok(self
+            .0
+            .sequencer_client
+            .send_transaction(NSSATransaction::Public(tx))
+            .await?)
     }
 
     pub async fn register_account(