From 8bc6256493564297cc53308431f2f894ad218e5f Mon Sep 17 00:00:00 2001 From: Thibault Martinez Date: Mon, 6 Nov 2023 16:19:14 +0100 Subject: [PATCH] Iterators experiment --- bindings/core/src/method_handler/wallet.rs | 4 +- cli/src/wallet_cli/mod.rs | 10 +- .../consolidate_outputs.rs | 4 +- sdk/examples/wallet/spammer.rs | 1 - sdk/src/wallet/core/mod.rs | 217 ++++++++++-------- .../wallet/operations/participation/mod.rs | 15 +- .../transaction/high_level/send_nft.rs | 9 +- .../operations/transaction/prepare_output.rs | 2 +- sdk/tests/wallet/consolidation.rs | 4 +- sdk/tests/wallet/output_preparation.rs | 2 +- 10 files changed, 144 insertions(+), 124 deletions(-) diff --git a/bindings/core/src/method_handler/wallet.rs b/bindings/core/src/method_handler/wallet.rs index 50948dadb7..8931fad35e 100644 --- a/bindings/core/src/method_handler/wallet.rs +++ b/bindings/core/src/method_handler/wallet.rs @@ -216,7 +216,7 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM } WalletMethod::Outputs { filter_options } => { let outputs = wallet.outputs(filter_options).await; - Response::OutputsData(outputs.iter().map(OutputDataDto::from).collect()) + Response::OutputsData(outputs.map(OutputDataDto::from).collect()) } WalletMethod::PendingTransactions => { let transactions = wallet.pending_transactions().await; @@ -399,7 +399,7 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM } WalletMethod::UnspentOutputs { filter_options } => { let outputs = wallet.unspent_outputs(filter_options).await; - Response::OutputsData(outputs.iter().map(OutputDataDto::from).collect()) + Response::OutputsData(outputs.map(OutputDataDto::from).collect()) } }; Ok(response) diff --git a/cli/src/wallet_cli/mod.rs b/cli/src/wallet_cli/mod.rs index f78c3dc5ce..9d6fb452dc 100644 --- a/cli/src/wallet_cli/mod.rs +++ b/cli/src/wallet_cli/mod.rs @@ -661,7 +661,7 @@ pub async fn output_command(wallet: &Wallet, selector: OutputSelector) -> Result let output = match selector { OutputSelector::Id(id) => wallet.get_output(&id).await, OutputSelector::Index(index) => { - let mut outputs = wallet.outputs(None).await; + let mut outputs = wallet.outputs(None).await.cloned().collect::>(); outputs.sort_unstable_by(outputs_ordering); outputs.into_iter().nth(index) } @@ -678,7 +678,7 @@ pub async fn output_command(wallet: &Wallet, selector: OutputSelector) -> Result /// `outputs` command pub async fn outputs_command(wallet: &Wallet) -> Result<(), Error> { - print_outputs(wallet.outputs(None).await, "Outputs:").await + print_outputs(wallet.outputs(None).await.cloned().collect(), "Outputs:").await } // `send` command @@ -835,7 +835,11 @@ pub async fn transactions_command(wallet: &Wallet, show_details: bool) -> Result /// `unspent-outputs` command pub async fn unspent_outputs_command(wallet: &Wallet) -> Result<(), Error> { - print_outputs(wallet.unspent_outputs(None).await, "Unspent outputs:").await + print_outputs( + wallet.unspent_outputs(None).await.cloned().collect(), + "Unspent outputs:", + ) + .await } pub async fn vote_command(wallet: &Wallet, event_id: ParticipationEventId, answers: Vec) -> Result<(), Error> { diff --git a/sdk/examples/how_tos/accounts_and_addresses/consolidate_outputs.rs b/sdk/examples/how_tos/accounts_and_addresses/consolidate_outputs.rs index bb08f13690..2cfb3100d6 100644 --- a/sdk/examples/how_tos/accounts_and_addresses/consolidate_outputs.rs +++ b/sdk/examples/how_tos/accounts_and_addresses/consolidate_outputs.rs @@ -44,7 +44,7 @@ async fn main() -> Result<()> { // output. let outputs = wallet.unspent_outputs(None).await; println!("Outputs BEFORE consolidation:"); - outputs.iter().enumerate().for_each(|(i, output_data)| { + outputs.enumerate().for_each(|(i, output_data)| { println!("OUTPUT #{i}"); println!( "- address: {:?}\n- amount: {:?}\n- native tokens: {:?}", @@ -80,7 +80,7 @@ async fn main() -> Result<()> { // Outputs after consolidation let outputs = wallet.unspent_outputs(None).await; println!("Outputs AFTER consolidation:"); - outputs.iter().enumerate().for_each(|(i, output_data)| { + outputs.enumerate().for_each(|(i, output_data)| { println!("OUTPUT #{i}"); println!( "- address: {:?}\n- amount: {:?}\n- native tokens: {:?}", diff --git a/sdk/examples/wallet/spammer.rs b/sdk/examples/wallet/spammer.rs index 9387836fd1..15bfb343e8 100644 --- a/sdk/examples/wallet/spammer.rs +++ b/sdk/examples/wallet/spammer.rs @@ -75,7 +75,6 @@ async fn main() -> Result<()> { ..Default::default() }) .await - .iter() .filter(|data| data.output.amount() >= SEND_AMOUNT) .count(); diff --git a/sdk/src/wallet/core/mod.rs b/sdk/src/wallet/core/mod.rs index f8a83acea3..ed3b0dee09 100644 --- a/sdk/src/wallet/core/mod.rs +++ b/sdk/src/wallet/core/mod.rs @@ -14,7 +14,7 @@ use crypto::keys::{ bip44::Bip44, }; use serde::{Deserialize, Serialize}; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard}; pub use self::builder::WalletBuilder; use super::types::{TransactionWithMetadata, TransactionWithMetadataDto}; @@ -103,6 +103,37 @@ pub struct WalletInner { pub(crate) storage_manager: tokio::sync::RwLock, } +struct OutputIterator<'a, T: Iterator> { + lock: RwLockReadGuard<'a, HashMap>, + iter: Option, +} + +// impl<'a, T: Iterator> OutputIterator<'a, T> { +// pub fn new(lock: RwLockReadGuard<'a, HashMap>) -> Self { +// Self { lock, iter: None } +// } +// } + +impl<'a> OutputIterator<'a, std::collections::hash_map::Values<'a, OutputId, OutputData>> { + pub fn new(lock: RwLockReadGuard<'a, HashMap>) -> Self { + Self { lock, iter: None } + } +} + +impl<'a> OutputIterator<'a, std::collections::hash_map::Values<'a, OutputId, OutputData>> { + pub fn set_iter(&'a mut self) { + self.iter = Some(self.lock.values()) + } +} + +impl<'a, T: Iterator> Iterator for OutputIterator<'a, T> { + type Item = &'a OutputData; + + fn next(&mut self) -> Option { + self.iter.as_mut().and_then(|i| i.next()) + } +} + /// Wallet data. #[derive(Clone, Debug, Eq, PartialEq)] pub struct WalletData { @@ -237,11 +268,11 @@ where self.inner.emit(wallet_event).await } - pub(crate) async fn data(&self) -> tokio::sync::RwLockReadGuard<'_, WalletData> { + pub(crate) async fn data(&self) -> RwLockReadGuard<'_, WalletData> { self.data.read().await } - pub(crate) async fn data_mut(&self) -> tokio::sync::RwLockWriteGuard<'_, WalletData> { + pub(crate) async fn data_mut(&self) -> RwLockWriteGuard<'_, WalletData> { self.data.write().await } @@ -295,164 +326,162 @@ where self.data().await.incoming_transactions.get(transaction_id).cloned() } - fn filter_outputs<'a>( - &self, - outputs: impl Iterator, - filter: impl Into>, - ) -> Vec { + fn filter_outputs<'a>(output: &OutputData, filter: impl Into>) -> bool { let filter = filter.into(); if let Some(filter) = filter { - let mut filtered_outputs = Vec::new(); - - for output in outputs { - match &output.output { - Output::Account(account) => { - if let Some(account_ids) = &filter.account_ids { - let account_id = account.account_id_non_null(&output.output_id); - if account_ids.contains(&account_id) { - filtered_outputs.push(output.clone()); - continue; - } + match &output.output { + Output::Account(account) => { + if let Some(account_ids) = &filter.account_ids { + let account_id = account.account_id_non_null(&output.output_id); + if account_ids.contains(&account_id) { + return true; } } - Output::Anchor(anchor) => { - if let Some(anchor_ids) = &filter.anchor_ids { - let anchor_id = anchor.anchor_id_non_null(&output.output_id); - if anchor_ids.contains(&anchor_id) { - filtered_outputs.push(output.clone()); - continue; - } + } + Output::Anchor(anchor) => { + if let Some(anchor_ids) = &filter.anchor_ids { + let anchor_id = anchor.anchor_id_non_null(&output.output_id); + if anchor_ids.contains(&anchor_id) { + return true; } } - Output::Foundry(foundry) => { - if let Some(foundry_ids) = &filter.foundry_ids { - let foundry_id = foundry.id(); - if foundry_ids.contains(&foundry_id) { - filtered_outputs.push(output.clone()); - continue; - } + } + Output::Foundry(foundry) => { + if let Some(foundry_ids) = &filter.foundry_ids { + let foundry_id = foundry.id(); + if foundry_ids.contains(&foundry_id) { + return true; } } - Output::Nft(nft) => { - if let Some(nft_ids) = &filter.nft_ids { - let nft_id = nft.nft_id_non_null(&output.output_id); - if nft_ids.contains(&nft_id) { - filtered_outputs.push(output.clone()); - continue; - } + } + Output::Nft(nft) => { + if let Some(nft_ids) = &filter.nft_ids { + let nft_id = nft.nft_id_non_null(&output.output_id); + if nft_ids.contains(&nft_id) { + return true; } } - Output::Delegation(delegation) => { - if let Some(delegation_ids) = &filter.delegation_ids { - let delegation_id = delegation.delegation_id_non_null(&output.output_id); - if delegation_ids.contains(&delegation_id) { - filtered_outputs.push(output.clone()); - continue; - } + } + Output::Delegation(delegation) => { + if let Some(delegation_ids) = &filter.delegation_ids { + let delegation_id = delegation.delegation_id_non_null(&output.output_id); + if delegation_ids.contains(&delegation_id) { + return true; } } - _ => {} } + _ => {} + } - // TODO filter based on slot index - // if let Some(lower_bound_booked_timestamp) = filter.lower_bound_booked_timestamp { - // if output.metadata.milestone_timestamp_booked() < lower_bound_booked_timestamp { - // continue; - // } - // } - // if let Some(upper_bound_booked_timestamp) = filter.upper_bound_booked_timestamp { - // if output.metadata.milestone_timestamp_booked() > upper_bound_booked_timestamp { - // continue; - // } - // } - - if let Some(output_types) = &filter.output_types { - if !output_types.contains(&output.output.kind()) { - continue; - } + // TODO filter based on slot index + // if let Some(lower_bound_booked_timestamp) = filter.lower_bound_booked_timestamp { + // if output.metadata.milestone_timestamp_booked() < lower_bound_booked_timestamp { + // continue; + // } + // } + // if let Some(upper_bound_booked_timestamp) = filter.upper_bound_booked_timestamp { + // if output.metadata.milestone_timestamp_booked() > upper_bound_booked_timestamp { + // continue; + // } + // } + + if let Some(output_types) = &filter.output_types { + if !output_types.contains(&output.output.kind()) { + return false; } + } - // Include the output if we're not filtering by IDs. - if filter.account_ids.is_none() - && filter.anchor_ids.is_none() - && filter.foundry_ids.is_none() - && filter.nft_ids.is_none() - && filter.delegation_ids.is_none() - { - filtered_outputs.push(output.clone()); - } + // Include the output if we're not filtering by IDs. + if filter.account_ids.is_none() + && filter.anchor_ids.is_none() + && filter.foundry_ids.is_none() + && filter.nft_ids.is_none() + && filter.delegation_ids.is_none() + { + return false; } - filtered_outputs + true } else { - outputs.cloned().collect() + true } } /// Returns outputs of the wallet. - pub async fn outputs(&self, filter: impl Into> + Send) -> Vec { - self.filter_outputs(self.data().await.outputs.values(), filter) + pub async fn outputs<'a>( + &'a self, + filter: impl Into> + Send, + ) -> impl Iterator { + let lock = RwLockReadGuard::map(self.data().await, |data| &data.outputs); + + let mut iterator = OutputIterator::new(lock); + iterator.set_iter(); + + iterator.filter(|output| output.output.is_account()) } /// Returns unspent outputs of the wallet. - pub async fn unspent_outputs(&self, filter: impl Into> + Send) -> Vec { - self.filter_outputs(self.data().await.unspent_outputs.values(), filter) + pub async fn unspent_outputs<'a>( + &'a self, + filter: impl Into> + Send, + ) -> impl Iterator { + let lock = RwLockReadGuard::map(self.data().await, |data| &data.unspent_outputs); + + let mut iterator = OutputIterator::new(lock); + // iterator.set_iter(); + + iterator.filter(|output| output.output.is_account()) } /// Gets the unspent account output matching the given ID. - pub async fn unspent_account_output(&self, account_id: &AccountId) -> Option { + pub async fn unspent_account_output(&self, account_id: &AccountId) -> Option<&OutputData> { self.unspent_outputs(FilterOptions { account_ids: Some([*account_id].into()), ..Default::default() }) .await - .first() - .cloned() + .next() } /// Gets the unspent anchor output matching the given ID. - pub async fn unspent_anchor_output(&self, anchor_id: &AnchorId) -> Option { + pub async fn unspent_anchor_output(&self, anchor_id: &AnchorId) -> Option<&OutputData> { self.unspent_outputs(FilterOptions { anchor_ids: Some([*anchor_id].into()), ..Default::default() }) .await - .first() - .cloned() + .next() } /// Gets the unspent foundry output matching the given ID. - pub async fn unspent_foundry_output(&self, foundry_id: &FoundryId) -> Option { + pub async fn unspent_foundry_output(&self, foundry_id: &FoundryId) -> Option<&OutputData> { self.unspent_outputs(FilterOptions { foundry_ids: Some([*foundry_id].into()), ..Default::default() }) .await - .first() - .cloned() + .next() } /// Gets the unspent nft output matching the given ID. - pub async fn unspent_nft_output(&self, nft_id: &NftId) -> Option { + pub async fn unspent_nft_output(&self, nft_id: &NftId) -> Option<&OutputData> { self.unspent_outputs(FilterOptions { nft_ids: Some([*nft_id].into()), ..Default::default() }) .await - .first() - .cloned() + .next() } /// Gets the unspent delegation output matching the given ID. - pub async fn unspent_delegation_output(&self, delegation_id: &DelegationId) -> Option { + pub async fn unspent_delegation_output(&self, delegation_id: &DelegationId) -> Option<&OutputData> { self.unspent_outputs(FilterOptions { delegation_ids: Some([*delegation_id].into()), ..Default::default() }) .await - .first() - .cloned() + .next() } /// Returns implicit accounts of the wallet. diff --git a/sdk/src/wallet/operations/participation/mod.rs b/sdk/src/wallet/operations/participation/mod.rs index cd7154580d..0984ea945c 100644 --- a/sdk/src/wallet/operations/participation/mod.rs +++ b/sdk/src/wallet/operations/participation/mod.rs @@ -72,16 +72,13 @@ where restored_spent_cached_outputs_len ); let outputs = self.outputs(None).await; - let participation_outputs = outputs - .into_iter() - .filter(|output_data| { - is_valid_participation_output(&output_data.output) + let participation_outputs = outputs.filter(|output_data| { + is_valid_participation_output(&output_data.output) // Check that the metadata exists, because otherwise we aren't participating for anything && output_data.output.features().and_then(|f| f.metadata()).is_some() // Don't add spent cached outputs, we have their data already and it can't change anymore && !spent_cached_outputs.contains_key(&output_data.output_id) - }) - .collect::>(); + }); let mut events = HashMap::new(); let mut spent_outputs = HashSet::new(); @@ -226,15 +223,13 @@ where /// Returns the voting output ("PARTICIPATION" tag). /// /// If multiple outputs with this tag exist, the one with the largest amount will be returned. - pub async fn get_voting_output(&self) -> Result> { + pub async fn get_voting_output(&self) -> Result> { log::debug!("[get_voting_output]"); Ok(self .unspent_outputs(None) .await - .iter() .filter(|output_data| is_valid_participation_output(&output_data.output)) - .max_by_key(|output_data| output_data.output.amount()) - .cloned()) + .max_by_key(|output_data| output_data.output.amount())) } /// Gets client for an event. diff --git a/sdk/src/wallet/operations/transaction/high_level/send_nft.rs b/sdk/src/wallet/operations/transaction/high_level/send_nft.rs index 11e045c813..73d1050bb9 100644 --- a/sdk/src/wallet/operations/transaction/high_level/send_nft.rs +++ b/sdk/src/wallet/operations/transaction/high_level/send_nft.rs @@ -95,7 +95,6 @@ where { log::debug!("[TRANSACTION] prepare_send_nft"); - let unspent_outputs = self.unspent_outputs(None).await; let token_supply = self.client().get_token_supply().await?; let mut outputs = Vec::new(); @@ -104,13 +103,7 @@ where self.client().bech32_hrp_matches(address.hrp()).await?; // Find nft output from the inputs - if let Some(nft_output_data) = unspent_outputs.iter().find(|o| { - if let Output::Nft(nft_output) = &o.output { - nft_id == nft_output.nft_id_non_null(&o.output_id) - } else { - false - } - }) { + if let Some(nft_output_data) = self.unspent_nft_output(&nft_id).await { if let Output::Nft(nft_output) = &nft_output_data.output { // Set the nft id and new address unlock condition let nft_builder = NftOutputBuilder::from(nft_output) diff --git a/sdk/src/wallet/operations/transaction/prepare_output.rs b/sdk/src/wallet/operations/transaction/prepare_output.rs index 3c0507d4c5..63c2efeeb7 100644 --- a/sdk/src/wallet/operations/transaction/prepare_output.rs +++ b/sdk/src/wallet/operations/transaction/prepare_output.rs @@ -277,7 +277,7 @@ where // Set new address unlock condition first_output_builder = first_output_builder.add_unlock_condition(AddressUnlockCondition::new(recipient_address)); - Ok((first_output_builder, existing_nft_output_data)) + Ok((first_output_builder, existing_nft_output_data.cloned())) } // Get a remainder address based on transaction_options or use the first account address diff --git a/sdk/tests/wallet/consolidation.rs b/sdk/tests/wallet/consolidation.rs index 64b92e0753..5bd6878e89 100644 --- a/sdk/tests/wallet/consolidation.rs +++ b/sdk/tests/wallet/consolidation.rs @@ -31,7 +31,7 @@ async fn consolidation() -> Result<()> { let balance = wallet_1.sync(None).await.unwrap(); assert_eq!(balance.base_coin().available(), 10 * amount); - assert_eq!(wallet_1.unspent_outputs(None).await.len(), 10); + assert_eq!(wallet_1.unspent_outputs(None).await.count(), 10); let tx = wallet_1 .consolidate_outputs(ConsolidationParams::new().with_force(true)) @@ -44,7 +44,7 @@ async fn consolidation() -> Result<()> { // Balance still the same assert_eq!(balance.base_coin().available(), 10 * amount); // Only one unspent output - assert_eq!(wallet_1.unspent_outputs(None).await.len(), 1); + assert_eq!(wallet_1.unspent_outputs(None).await.count(), 1); tear_down(storage_path_0)?; tear_down(storage_path_1)?; diff --git a/sdk/tests/wallet/output_preparation.rs b/sdk/tests/wallet/output_preparation.rs index 38063b95db..8e9990a1ce 100644 --- a/sdk/tests/wallet/output_preparation.rs +++ b/sdk/tests/wallet/output_preparation.rs @@ -766,7 +766,7 @@ async fn prepare_output_only_single_nft() -> Result<()> { let balance = wallet_1.sync(None).await?; assert_eq!(balance.nfts().len(), 1); - let nft_data = &wallet_1.unspent_outputs(None).await[0]; + let nft_data = &wallet_1.unspent_outputs(None).await.next().unwrap(); let nft_id = *balance.nfts().first().unwrap(); // Send NFT back to first wallet let output = wallet_1