From 7d5e1492c46dc196c41ea3e8e046ff46dc2182e7 Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Mon, 4 May 2026 20:12:54 -0300
Subject: [PATCH] update insert_private_account_data to take account kind
 instead of identifier

---
 wallet/src/chain_storage.rs | 13 ++++++++++---
 wallet/src/lib.rs           | 16 ++++++++--------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/wallet/src/chain_storage.rs b/wallet/src/chain_storage.rs
index 8d168d8e..4b160b66 100644
--- a/wallet/src/chain_storage.rs
+++ b/wallet/src/chain_storage.rs
@@ -11,6 +11,7 @@ use key_protocol::{
 };
 use log::debug;
 use nssa::program::Program;
+use nssa_core::PrivateAccountKind;
 
 use crate::config::{InitialAccountData, Label, PersistentAccountData, WalletConfig};
 
@@ -190,10 +191,11 @@ impl WalletChainStore {
     pub fn insert_private_account_data(
         &mut self,
         account_id: nssa::AccountId,
-        identifier: nssa_core::Identifier,
+        kind: &PrivateAccountKind,
         account: nssa_core::account::Account,
     ) {
         debug!("inserting at address {account_id}, this account {account:?}");
+        let identifier = kind.identifier();
 
         // Update default accounts if present
         if let Entry::Occupied(mut entry) = self
@@ -237,8 +239,13 @@ impl WalletChainStore {
         } else {
             // Node not yet in account_id_map — find it by checking all nodes
             for (ci, node) in &mut self.user_data.private_key_tree.key_map {
-                let expected_id =
-                    nssa::AccountId::from((&node.value.0.nullifier_public_key, identifier));
+                let npk = &node.value.0.nullifier_public_key;
+                let expected_id = match kind {
+                    PrivateAccountKind::Account(id) => nssa::AccountId::from((npk, *id)),
+                    PrivateAccountKind::Pda { program_id, seed, identifier: id } => {
+                        nssa::AccountId::for_private_pda(program_id, seed, npk, *id)
+                    }
+                };
                 if expected_id == account_id {
                     if let Some((_, acc)) =
                         node.value.1.iter_mut().find(|(id, _)| *id == identifier)
diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs
index 31e283a9..ee9f3dd1 100644
--- a/wallet/src/lib.rs
+++ b/wallet/src/lib.rs
@@ -284,7 +284,7 @@ impl WalletCore {
             .nullifier_public_key;
         let account_id = AccountId::from((&npk, identifier));
         self.storage
-            .insert_private_account_data(account_id, identifier, Account::default());
+            .insert_private_account_data(account_id, &PrivateAccountKind::Account(identifier), Account::default());
         (account_id, cci)
     }
 
@@ -371,7 +371,7 @@ impl WalletCore {
                     println!("Received new acc {res_acc:#?}");
 
                     self.storage
-                        .insert_private_account_data(*acc_account_id, kind.identifier(), res_acc);
+                        .insert_private_account_data(*acc_account_id, &kind, res_acc);
                 }
                 AccDecodeData::Skip => {}
             }
@@ -547,27 +547,27 @@ impl WalletCore {
                         )
                         .map(|(kind, res_acc)| {
                             let npk = &key_chain.nullifier_public_key;
-                            let (account_id, identifier) = match kind {
+                            let account_id = match &kind {
                                 PrivateAccountKind::Account(identifier) => {
-                                    (nssa::AccountId::from((npk, identifier)), identifier)
+                                    nssa::AccountId::from((npk, *identifier))
                                 }
                                 PrivateAccountKind::Pda { program_id, seed, identifier } => {
-                                    (nssa::AccountId::for_private_pda(&program_id, &seed, npk, identifier), identifier)
+                                    nssa::AccountId::for_private_pda(program_id, seed, npk, *identifier)
                                 }
                             };
-                            (account_id, identifier, res_acc)
+                            (account_id, kind, res_acc)
                         })
                     })
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
 
-        for (affected_account_id, identifier, new_acc) in affected_accounts {
+        for (affected_account_id, kind, new_acc) in affected_accounts {
             info!(
                 "Received new account for account_id {affected_account_id:#?} with account object {new_acc:#?}"
             );
             self.storage
-                .insert_private_account_data(affected_account_id, identifier, new_acc);
+                .insert_private_account_data(affected_account_id, &kind, new_acc);
         }
     }