From e19c9ff20ae9aea139c037d8f3a9c422028ac94f Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Fri, 24 Apr 2026 00:04:22 -0300 Subject: [PATCH] return impl iterator --- .../key_management/key_tree/keys_private.rs | 3 +-- .../src/key_management/key_tree/keys_public.rs | 4 ++-- .../src/key_management/key_tree/mod.rs | 8 ++++---- .../src/key_management/key_tree/traits.rs | 2 +- wallet/src/chain_storage.rs | 18 ++++++++---------- wallet/src/helperfunctions.rs | 6 +++--- wallet/src/lib.rs | 2 +- 7 files changed, 20 insertions(+), 23 deletions(-) diff --git a/key_protocol/src/key_management/key_tree/keys_private.rs b/key_protocol/src/key_management/key_tree/keys_private.rs index ee6584f2..459fdc23 100644 --- a/key_protocol/src/key_management/key_tree/keys_private.rs +++ b/key_protocol/src/key_management/key_tree/keys_private.rs @@ -134,14 +134,13 @@ impl KeyTreeNode for ChildKeysPrivate { self.nth_child(cci) } - fn account_ids(&self) -> Vec { + fn account_ids(&self) -> impl Iterator { self.value .1 .iter() .map(|(identifier, _)| { nssa::AccountId::from((&self.value.0.nullifier_public_key, *identifier)) }) - .collect() } } diff --git a/key_protocol/src/key_management/key_tree/keys_public.rs b/key_protocol/src/key_management/key_tree/keys_public.rs index 63678e4c..3ab9cc35 100644 --- a/key_protocol/src/key_management/key_tree/keys_public.rs +++ b/key_protocol/src/key_management/key_tree/keys_public.rs @@ -115,8 +115,8 @@ impl KeyTreeNode for ChildKeysPublic { self.nth_child(cci) } - fn account_ids(&self) -> Vec { - vec![self.account_id()] + fn account_ids(&self) -> impl Iterator { + std::iter::once(self.account_id()) } } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 002d1746..0702293d 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -208,7 +208,7 @@ impl KeyTree { ) -> Option<(nssa::AccountId, ChainIndex)> { let cci = self.generate_new_node(parent_cci)?; let node = self.key_map.get(&cci)?; - let account_id = *node.account_ids().first()?; + let account_id = node.account_ids().next()?; Some((account_id, cci)) } @@ -217,7 +217,7 @@ impl KeyTree { pub fn generate_new_public_node_layered(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { let cci = self.generate_new_node_layered()?; let node = self.key_map.get(&cci)?; - let account_id = *node.account_ids().first()?; + let account_id = node.account_ids().next()?; Some((account_id, cci)) } @@ -298,9 +298,9 @@ impl KeyTree { 'outer: for i in (1..depth).rev() { println!("Cleanup of tree at depth {i}"); for id in ChainIndex::chain_ids_at_depth(i) { - if let Some(node) = self.key_map.get(&id) { + if let Some(node) = self.key_map.get(&id).cloned() { if node.value.1.is_empty() { - let account_ids: Vec<_> = node.account_ids(); + let account_ids = node.account_ids(); self.key_map.remove(&id); for addr in account_ids { self.account_id_map.remove(&addr); diff --git a/key_protocol/src/key_management/key_tree/traits.rs b/key_protocol/src/key_management/key_tree/traits.rs index 859070fd..71ca4743 100644 --- a/key_protocol/src/key_management/key_tree/traits.rs +++ b/key_protocol/src/key_management/key_tree/traits.rs @@ -4,5 +4,5 @@ pub trait KeyTreeNode: Sized { #[must_use] fn derive_child(&self, cci: u32) -> Self; #[must_use] - fn account_ids(&self) -> Vec; + fn account_ids(&self) -> impl Iterator; } diff --git a/wallet/src/chain_storage.rs b/wallet/src/chain_storage.rs index d5534d01..903c5b3b 100644 --- a/wallet/src/chain_storage.rs +++ b/wallet/src/chain_storage.rs @@ -7,7 +7,7 @@ use key_protocol::{ key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex}, secret_holders::SeedHolder, }, - key_protocol_core::NSSAUserData, + key_protocol_core::{NSSAUserData, UserPrivateAccountData}, }; use log::debug; use nssa::program::Program; @@ -79,8 +79,7 @@ impl WalletChainStore { InitialAccountData::Private(data) => { private_init_acc_map.insert( data.account_id, - (data.key_chain, vec![(data.identifier, data.account)]), - ); + UserPrivateAccountData{ key_chain: data.key_chain, accounts: vec![(data.identifier, data.account)]}); } }, } @@ -121,8 +120,7 @@ impl WalletChainStore { account.program_owner = Program::authenticated_transfer_program().id(); private_init_acc_map.insert( data.account_id, - (data.key_chain, vec![(data.identifier, account)]), - ); + UserPrivateAccountData{key_chain: data.key_chain, accounts:vec![(data.identifier, account)]}); } } } @@ -186,19 +184,19 @@ impl WalletChainStore { .default_user_private_accounts .entry(account_id) { - let (key_chain, entries) = entry.get_mut(); - let identifier = entries + let entry = entry.get_mut(); + let identifier = entry.accounts .iter() .find_map(|(id, _)| { - (nssa::AccountId::from((&key_chain.nullifier_public_key, *id)) == account_id) + (nssa::AccountId::from((&entry.key_chain.nullifier_public_key, *id)) == account_id) .then_some(*id) }) .unwrap_or(0); // Update existing entry or insert new one - if let Some((_, acc)) = entries.iter_mut().find(|(id, _)| *id == identifier) { + if let Some((_, acc)) = entry.accounts.iter_mut().find(|(id, _)| *id == identifier) { *acc = account; } else { - entries.push((identifier, account)); + entry.accounts.push((identifier, account)); } return; } diff --git a/wallet/src/helperfunctions.rs b/wallet/src/helperfunctions.rs index 26141a7d..bdefe10a 100644 --- a/wallet/src/helperfunctions.rs +++ b/wallet/src/helperfunctions.rs @@ -212,13 +212,13 @@ pub fn produce_data_for_storage( ); } - for (account_id, (key_chain, entries)) in &user_data.default_user_private_accounts { - for (identifier, account) in entries { + for (account_id, entry) in &user_data.default_user_private_accounts { + for (identifier, account) in &entry.accounts { vec_for_storage.push( InitialAccountData::Private(Box::new(PrivateAccountPrivateInitialData { account_id: *account_id, account: account.clone(), - key_chain: key_chain.clone(), + key_chain: entry.key_chain.clone(), identifier: *identifier, })) .into(), diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 62b40b61..155b4026 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -508,7 +508,7 @@ impl WalletCore { .user_data .default_user_private_accounts .values() - .map(|(key_chain, _)| (key_chain, None)) + .map(|entry| (&entry.key_chain, None)) .chain( self.storage .user_data