Allow some wallet data to be accessible by reference (#1582)

* allow some wallet data to be accessible by reference * PR suggestions --------- Co-authored-by: Thibault Martinez <[email protected]>
iotaledger · Nov 30, 2023 · 1f6ff33 · 1f6ff33
1 parent d58d964
commit 1f6ff33
Show file tree

Hide file tree

Showing 13 changed files with 357 additions and 346 deletions.
diff --git a/bindings/core/src/method_handler/wallet.rs b/bindings/core/src/method_handler/wallet.rs
@@ -17,8 +17,7 @@ use crate::{method::WalletMethod, response::Response};
 pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletMethod) -> crate::Result<Response> {
     let response = match method {
         WalletMethod::Accounts => {
-            let accounts = wallet.accounts().await;
-            Response::OutputsData(accounts.iter().map(OutputDataDto::from).collect())
+            Response::OutputsData(wallet.data().await.accounts().map(OutputDataDto::from).collect())
         }
         #[cfg(feature = "stronghold")]
         WalletMethod::Backup { destination, password } => {
@@ -154,18 +153,22 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM
             let output = wallet.get_foundry_output(token_id).await?;
             Response::Output(output)
         }
-        WalletMethod::GetIncomingTransaction { transaction_id } => {
-            let transaction = wallet.get_incoming_transaction(&transaction_id).await;
-
-            transaction.map_or_else(
+        WalletMethod::GetIncomingTransaction { transaction_id } => wallet
+            .data()
+            .await
+            .get_incoming_transaction(&transaction_id)
+            .map_or_else(
                 || Response::Transaction(None),
-                |transaction| Response::Transaction(Some(Box::new(TransactionWithMetadataDto::from(&transaction)))),
-            )
-        }
-        WalletMethod::GetOutput { output_id } => {
-            let output_data = wallet.get_output(&output_id).await;
-            Response::OutputData(output_data.as_ref().map(OutputDataDto::from).map(Box::new))
-        }
+                |transaction| Response::Transaction(Some(Box::new(TransactionWithMetadataDto::from(transaction)))),
+            ),
+        WalletMethod::GetOutput { output_id } => Response::OutputData(
+            wallet
+                .data()
+                .await
+                .get_output(&output_id)
+                .map(OutputDataDto::from)
+                .map(Box::new),
+        ),
         #[cfg(feature = "participation")]
         WalletMethod::GetParticipationEvent { event_id } => {
             let event_and_nodes = wallet.get_participation_event(event_id).await?;
@@ -191,10 +194,14 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM
             let overview = wallet.get_participation_overview(event_ids).await?;
             Response::ParticipationOverview(overview)
         }
-        WalletMethod::GetTransaction { transaction_id } => {
-            let transaction = wallet.get_transaction(&transaction_id).await;
-            Response::Transaction(transaction.as_ref().map(TransactionWithMetadataDto::from).map(Box::new))
-        }
+        WalletMethod::GetTransaction { transaction_id } => Response::Transaction(
+            wallet
+                .data()
+                .await
+                .get_transaction(&transaction_id)
+                .map(TransactionWithMetadataDto::from)
+                .map(Box::new),
+        ),
         #[cfg(feature = "participation")]
         WalletMethod::GetVotingPower => {
             let voting_power = wallet.get_voting_power().await?;
@@ -204,26 +211,43 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM
             let implicit_account_creation_address = wallet.implicit_account_creation_address().await?;
             Response::Bech32Address(implicit_account_creation_address)
         }
+        WalletMethod::ImplicitAccounts => Response::OutputsData(
+            wallet
+                .data()
+                .await
+                .implicit_accounts()
+                .map(OutputDataDto::from)
+                .collect(),
+        ),
         WalletMethod::PrepareImplicitAccountTransition { output_id } => {
             let data = wallet.prepare_implicit_account_transition(&output_id).await?;
             Response::PreparedTransaction(PreparedTransactionDataDto::from(&data))
         }
-        WalletMethod::ImplicitAccounts => {
-            let implicit_accounts = wallet.implicit_accounts().await;
-            Response::OutputsData(implicit_accounts.iter().map(OutputDataDto::from).collect())
-        }
-        WalletMethod::IncomingTransactions => {
-            let transactions = wallet.incoming_transactions().await;
-            Response::Transactions(transactions.iter().map(TransactionWithMetadataDto::from).collect())
-        }
+        WalletMethod::IncomingTransactions => Response::Transactions(
+            wallet
+                .data()
+                .await
+                .incoming_transactions()
+                .values()
+                .map(TransactionWithMetadataDto::from)
+                .collect(),
+        ),
         WalletMethod::Outputs { filter_options } => {
-            let outputs = wallet.outputs(filter_options).await;
-            Response::OutputsData(outputs.iter().map(OutputDataDto::from).collect())
-        }
-        WalletMethod::PendingTransactions => {
-            let transactions = wallet.pending_transactions().await;
-            Response::Transactions(transactions.iter().map(TransactionWithMetadataDto::from).collect())
-        }
+            let wallet_data = wallet.data().await;
+            Response::OutputsData(if let Some(filter) = filter_options {
+                wallet_data.filtered_outputs(filter).map(OutputDataDto::from).collect()
+            } else {
+                wallet_data.outputs().values().map(OutputDataDto::from).collect()
+            })
+        }
+        WalletMethod::PendingTransactions => Response::Transactions(
+            wallet
+                .data()
+                .await
+                .pending_transactions()
+                .map(TransactionWithMetadataDto::from)
+                .collect(),
+        ),
         WalletMethod::PrepareBurn { burn, options } => {
             let data = wallet.prepare_burn(burn, options).await?;
             Response::PreparedTransaction(PreparedTransactionDataDto::from(&data))
@@ -381,13 +405,29 @@ pub(crate) async fn call_wallet_method_internal(wallet: &Wallet, method: WalletM
             Response::SentTransaction(TransactionWithMetadataDto::from(&transaction))
         }
         WalletMethod::Sync { options } => Response::Balance(wallet.sync(options).await?),
-        WalletMethod::Transactions => {
-            let transactions = wallet.transactions().await;
-            Response::Transactions(transactions.iter().map(TransactionWithMetadataDto::from).collect())
-        }
+        WalletMethod::Transactions => Response::Transactions(
+            wallet
+                .data()
+                .await
+                .transactions()
+                .values()
+                .map(TransactionWithMetadataDto::from)
+                .collect(),
+        ),
         WalletMethod::UnspentOutputs { filter_options } => {
-            let outputs = wallet.unspent_outputs(filter_options).await;
-            Response::OutputsData(outputs.iter().map(OutputDataDto::from).collect())
+            let wallet_data = wallet.data().await;
+            Response::OutputsData(if let Some(filter) = filter_options {
+                wallet_data
+                    .filtered_unspent_outputs(filter)
+                    .map(OutputDataDto::from)
+                    .collect()
+            } else {
+                wallet_data
+                    .unspent_outputs()
+                    .values()
+                    .map(OutputDataDto::from)
+                    .collect()
+            })
         }
     };
     Ok(response)

diff --git a/cli/src/wallet_cli/mod.rs b/cli/src/wallet_cli/mod.rs
@@ -23,9 +23,8 @@ use iota_sdk::{
     },
     utils::ConvertTo,
     wallet::{
-        types::{OutputData, TransactionWithMetadata},
-        ConsolidationParams, CreateNativeTokenParams, MintNftParams, OutputsToClaim, SendNativeTokenParams,
-        SendNftParams, SendParams, SyncOptions, TransactionOptions, Wallet,
+        types::OutputData, ConsolidationParams, CreateNativeTokenParams, MintNftParams, OutputsToClaim,
+        SendNativeTokenParams, SendNftParams, SendParams, SyncOptions, TransactionOptions, Wallet,
     },
     U256,
 };
@@ -313,7 +312,7 @@ impl FromStr for OutputSelector {
 
 // `accounts` command
 pub async fn accounts_command(wallet: &Wallet) -> Result<(), Error> {
-    print_outputs(wallet.accounts().await, "Accounts:").await
+    print_outputs(wallet.data().await.accounts().cloned().collect(), "Accounts:")
 }
 
 // `address` command
@@ -413,9 +412,9 @@ pub async fn claimable_outputs_command(wallet: &Wallet) -> Result<(), Error> {
         .iter()
         .filter_map(|(output_id, unlockable)| unlockable.then_some(output_id))
     {
+        let wallet_data = wallet.data().await;
         // Unwrap: for the iterated `OutputId`s this call will always return `Some(...)`.
-        let output_data = wallet.get_output(output_id).await.unwrap();
-        let output = output_data.output;
+        let output = &wallet_data.get_output(output_id).unwrap().output;
         let kind = match output {
             Output::Nft(_) => "Nft",
             Output::Basic(_) => "Basic",
@@ -598,7 +597,10 @@ pub async fn implicit_account_transition_command(wallet: &Wallet, output_id: Out
 
 // `implicit-accounts` command
 pub async fn implicit_accounts_command(wallet: &Wallet) -> Result<(), Error> {
-    print_outputs(wallet.implicit_accounts().await, "Implicit accounts:").await
+    print_outputs(
+        wallet.data().await.implicit_accounts().cloned().collect(),
+        "Implicit accounts:",
+    )
 }
 
 // `melt-native-token` command
@@ -684,11 +686,12 @@ pub async fn node_info_command(wallet: &Wallet) -> Result<(), Error> {
 
 /// `output` command
 pub async fn output_command(wallet: &Wallet, selector: OutputSelector) -> Result<(), Error> {
+    let wallet_data = wallet.data().await;
     let output = match selector {
-        OutputSelector::Id(id) => wallet.get_output(&id).await,
+        OutputSelector::Id(id) => wallet_data.get_output(&id),
         OutputSelector::Index(index) => {
-            let mut outputs = wallet.outputs(None).await;
-            outputs.sort_unstable_by(outputs_ordering);
+            let mut outputs = wallet_data.outputs().values().collect::<Vec<_>>();
+            outputs.sort_unstable_by_key(|o| o.output_id);
             outputs.into_iter().nth(index)
         }
     };
@@ -704,7 +707,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.data().await.outputs().values().cloned().collect(), "Outputs:")
 }
 
 // `send` command
@@ -818,11 +821,12 @@ pub async fn sync_command(wallet: &Wallet) -> Result<(), Error> {
 
 /// `transaction` command
 pub async fn transaction_command(wallet: &Wallet, selector: TransactionSelector) -> Result<(), Error> {
-    let mut transactions = wallet.transactions().await;
+    let wallet_data = wallet.data().await;
     let transaction = match selector {
-        TransactionSelector::Id(id) => transactions.into_iter().find(|tx| tx.transaction_id == id),
+        TransactionSelector::Id(id) => wallet_data.get_transaction(&id),
         TransactionSelector::Index(index) => {
-            transactions.sort_unstable_by(transactions_ordering);
+            let mut transactions = wallet_data.transactions().values().collect::<Vec<_>>();
+            transactions.sort_unstable_by(|a, b| b.timestamp.cmp(&a.timestamp));
             transactions.into_iter().nth(index)
         }
     };
@@ -838,8 +842,9 @@ pub async fn transaction_command(wallet: &Wallet, selector: TransactionSelector)
 
 /// `transactions` command
 pub async fn transactions_command(wallet: &Wallet, show_details: bool) -> Result<(), Error> {
-    let mut transactions = wallet.transactions().await;
-    transactions.sort_unstable_by(transactions_ordering);
+    let wallet_data = wallet.data().await;
+    let mut transactions = wallet_data.transactions().values().collect::<Vec<_>>();
+    transactions.sort_unstable_by(|a, b| b.timestamp.cmp(&a.timestamp));
 
     if transactions.is_empty() {
         println_log_info!("No transactions found");
@@ -861,7 +866,10 @@ 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.data().await.unspent_outputs().values().cloned().collect(),
+        "Unspent outputs:",
+    )
 }
 
 pub async fn vote_command(wallet: &Wallet, event_id: ParticipationEventId, answers: Vec<u8>) -> Result<(), Error> {
@@ -950,7 +958,6 @@ async fn print_wallet_address(wallet: &Wallet) -> Result<(), Error> {
         address.inner()
     );
 
-    let unspent_outputs = wallet.unspent_outputs(None).await;
     let slot_index = wallet.client().get_slot_index().await?;
 
     let mut output_ids = Vec::new();
@@ -962,7 +969,7 @@ async fn print_wallet_address(wallet: &Wallet) -> Result<(), Error> {
     let mut delegations = Vec::new();
     let mut anchors = Vec::new();
 
-    for output_data in unspent_outputs {
+    for output_data in wallet.data().await.unspent_outputs().values() {
         let output_id = output_data.output_id;
         output_ids.push(output_id);
 
@@ -1216,12 +1223,12 @@ pub async fn prompt_internal(
     Ok(PromptResponse::Reprompt)
 }
 
-async fn print_outputs(mut outputs: Vec<OutputData>, title: &str) -> Result<(), Error> {
+fn print_outputs(mut outputs: Vec<OutputData>, title: &str) -> Result<(), Error> {
     if outputs.is_empty() {
         println_log_info!("No outputs found");
     } else {
         println_log_info!("{title}");
-        outputs.sort_unstable_by(outputs_ordering);
+        outputs.sort_unstable_by_key(|o| o.output_id);
 
         for (i, output_data) in outputs.into_iter().enumerate() {
             let kind_str = if output_data.output.is_implicit_account() {
@@ -1242,11 +1249,3 @@ async fn print_outputs(mut outputs: Vec<OutputData>, title: &str) -> Result<(),
 
     Ok(())
 }
-
-fn outputs_ordering(a: &OutputData, b: &OutputData) -> std::cmp::Ordering {
-    a.output_id.cmp(&b.output_id)
-}
-
-fn transactions_ordering(a: &TransactionWithMetadata, b: &TransactionWithMetadata) -> std::cmp::Ordering {
-    b.timestamp.cmp(&a.timestamp)
-}
diff --git a/sdk/examples/how_tos/accounts_and_addresses/consolidate_outputs.rs b/sdk/examples/how_tos/accounts_and_addresses/consolidate_outputs.rs
@@ -48,17 +48,22 @@ async fn main() -> Result<()> {
     // unlock condition and it is an `AddressUnlockCondition`, and so they are valid for consolidation. They have the
     // same `AddressUnlockCondition`(the address of the wallet), so they will be consolidated into one
     // output.
-    let outputs = wallet.unspent_outputs(None).await;
     println!("Outputs BEFORE consolidation:");
-    outputs.iter().enumerate().for_each(|(i, output_data)| {
-        println!("OUTPUT #{i}");
-        println!(
-            "- address: {:?}\n- amount: {:?}\n- native tokens: {:?}",
-            output_data.address.clone().to_bech32_unchecked("rms"),
-            output_data.output.amount(),
-            output_data.output.native_token()
-        )
-    });
+    wallet
+        .data()
+        .await
+        .unspent_outputs()
+        .values()
+        .enumerate()
+        .for_each(|(i, output_data)| {
+            println!("OUTPUT #{i}");
+            println!(
+                "- address: {:?}\n- amount: {:?}\n- native tokens: {:?}",
+                output_data.address.clone().to_bech32_unchecked("rms"),
+                output_data.output.amount(),
+                output_data.output.native_token()
+            )
+        });
 
     println!("Sending consolidation transaction...");
 
@@ -84,17 +89,22 @@ async fn main() -> Result<()> {
     println!("Wallet synced");
 
     // Outputs after consolidation
-    let outputs = wallet.unspent_outputs(None).await;
     println!("Outputs AFTER consolidation:");
-    outputs.iter().enumerate().for_each(|(i, output_data)| {
-        println!("OUTPUT #{i}");
-        println!(
-            "- address: {:?}\n- amount: {:?}\n- native tokens: {:?}",
-            output_data.address.clone().to_bech32_unchecked("rms"),
-            output_data.output.amount(),
-            output_data.output.native_token()
-        )
-    });
+    wallet
+        .data()
+        .await
+        .unspent_outputs()
+        .values()
+        .enumerate()
+        .for_each(|(i, output_data)| {
+            println!("OUTPUT #{i}");
+            println!(
+                "- address: {:?}\n- amount: {:?}\n- native tokens: {:?}",
+                output_data.address.clone().to_bech32_unchecked("rms"),
+                output_data.output.amount(),
+                output_data.output.native_token()
+            )
+        });
 
     Ok(())
 }
diff --git a/sdk/examples/how_tos/accounts_and_addresses/list_outputs.rs b/sdk/examples/how_tos/accounts_and_addresses/list_outputs.rs
@@ -30,13 +30,13 @@ async fn main() -> Result<()> {
 
     // Print output ids
     println!("Output ids:");
-    for output in wallet.outputs(None).await {
+    for output in wallet.data().await.outputs().values() {
         println!("{}", output.output_id);
     }
 
     // Print unspent output ids
     println!("Unspent output ids:");
-    for output in wallet.unspent_outputs(None).await {
+    for output in wallet.data().await.unspent_outputs().values() {
         println!("{}", output.output_id);
     }