refactor(wallet): use iterators and adaptors in preselect_utxos

There were multiple calls for de-duplication of selected UTxOs.

As the test `test_filter_duplicates` shows, there are four possible
cases for duplication of UTxOs while feeding the coin selection
algorithms.

1. no duplication: out of concern
2. duplication in the required utxos only: covered by the source of
   `required_utxos`, `Wallet::list_unspent`, which roots back the
   provided `UTxOs` to a `HashMap` which should avoid any duplication by
   definition
3. duplication in the optional utxos only: is the only one possible as
   optional `UTxOs` are stored in a `Vec` and no checks are performed
   about the duplicity of their members.
4. duplication across the required and optional utxos: is already
   covered by `Wallet::preselect_utxos`, which avoid the processing of
   required UTxOs while listing the unspent available UTxOs in the
   wallet.

This refactor changes:
- `TxParams::utxos` type to be `HashSet<LocalOutput>` avoiding the
  duplication case 3, and allowing a query closer to O(1) on avg. to
  cover duplication case 4 (before was O(n) where n is the size of
  required utxos).
- Moves the computation of the `WeightedUtxos` to the last part of UTxO
  filtering, allowing the unification of the computation for local
  outputs.
- Removes some extra allocations done for helpers structures or
  intermediate results while filtering UTxOs.
- Allows for future integration of UTxO filtering methods for other
  utilities.
- Adds more comments for each filtering step.

With these changes all four cases would be covered, and
`coin_selection::filter_duplicates` would be no longer needed.

Author: nymius

crates/wallet/src/wallet/mod.rs

@@ -36,7 +36,7 @@ use bdk_chain::{
 use bitcoin::{
     absolute,
     consensus::encode::serialize,
-    constants::{genesis_block, COINBASE_MATURITY},
+    constants::genesis_block,
     psbt,
     secp256k1::Secp256k1,
     sighash::{EcdsaSighashType, TapSighashType},
@@ -1417,8 +1417,33 @@ impl Wallet {
 
         fee_amount += fee_rate * tx.weight();
 
-        let (required_utxos, optional_utxos) =
-            self.preselect_utxos(&params, Some(current_height.to_consensus_u32()));
+        let (required_utxos, optional_utxos) = {
+            // NOTE: manual selection overrides unspendable
+            let mut required: Vec<WeightedUtxo> = params
+                .utxos
+                .clone()
+                .into_iter()
+                // transform required UTxOs into a WeightedUtxo
+                .map(|utxo| WeightedUtxo {
+                    satisfaction_weight: self
+                        .public_descriptor(utxo.keychain)
+                        .max_weight_to_satisfy()
+                        .unwrap(),
+                    utxo: Utxo::Local(utxo),
+                })
+                // include foreign UTxOs as required
+                .chain(params.foreign_utxos.clone())
+                .collect();
+            let optional = self.filter_utxos(&params, current_height.to_consensus_u32());
+
+            // if drain_wallet is true, all UTxOs are required
+            if params.drain_wallet {
+                required.extend(optional);
+                (required, vec![])
+            } else {
+                (required, optional)
+            }
+        };
 
         // get drain script
         let mut drain_index = Option::<(KeychainKind, u32)>::None;
@@ -1618,60 +1643,64 @@ impl Wallet {
             .map_err(|_| BuildFeeBumpError::FeeRateUnavailable)?;
 
         // remove the inputs from the tx and process them
-        let original_txin = tx.input.drain(..).collect::<Vec<_>>();
-        let original_utxos = original_txin
-            .iter()
-            .map(|txin| -> Result<_, BuildFeeBumpError> {
-                let prev_tx = graph
-                    .get_tx(txin.previous_output.txid)
-                    .ok_or(BuildFeeBumpError::UnknownUtxo(txin.previous_output))?;
-                let txout = &prev_tx.output[txin.previous_output.vout as usize];
-
-                let chain_position = chain_positions
-                    .get(&txin.previous_output.txid)
-                    .cloned()
-                    .ok_or(BuildFeeBumpError::UnknownUtxo(txin.previous_output))?;
-
-                let weighted_utxo = match txout_index.index_of_spk(txout.script_pubkey.clone()) {
-                    Some(&(keychain, derivation_index)) => {
-                        let satisfaction_weight = self
-                            .public_descriptor(keychain)
-                            .max_weight_to_satisfy()
-                            .unwrap();
-                        WeightedUtxo {
-                            utxo: Utxo::Local(LocalOutput {
-                                outpoint: txin.previous_output,
-                                txout: txout.clone(),
-                                keychain,
-                                is_spent: true,
-                                derivation_index,
-                                chain_position,
-                            }),
-                            satisfaction_weight,
-                        }
-                    }
-                    None => {
-                        let satisfaction_weight = Weight::from_wu_usize(
-                            serialize(&txin.script_sig).len() * 4 + serialize(&txin.witness).len(),
-                        );
-                        WeightedUtxo {
-                            utxo: Utxo::Foreign {
-                                outpoint: txin.previous_output,
-                                sequence: txin.sequence,
-                                psbt_input: Box::new(psbt::Input {
+        let (local, foreign): (HashSet<LocalOutput>, Vec<WeightedUtxo>) =
+            tx.input.drain(..).try_fold(
+                (HashSet::<LocalOutput>::new(), Vec::<WeightedUtxo>::new()),
+                |(mut local, mut foreign), txin| -> Result<_, BuildFeeBumpError> {
+                    let prev_tx = graph
+                        .get_tx(txin.previous_output.txid)
+                        .ok_or(BuildFeeBumpError::UnknownUtxo(txin.previous_output))?;
+                    let txout = &prev_tx.output[txin.previous_output.vout as usize];
+
+                    let chain_position =
+                        chain_positions
+                            .get(&txin.previous_output.txid)
+                            .cloned()
+                            .ok_or(BuildFeeBumpError::UnknownUtxo(txin.previous_output))?;
+
+                    match txout_index.index_of_spk(txout.script_pubkey.clone()) {
+                        Some(&(keychain, derivation_index)) => local.insert(LocalOutput {
+                            outpoint: txin.previous_output,
+                            txout: txout.clone(),
+                            keychain,
+                            is_spent: true,
+                            derivation_index,
+                            chain_position,
+                        }),
+                        None => {
+                            let satisfaction_weight = Weight::from_wu_usize(
+                                serialize(&txin.script_sig).len() * 4
+                                    + serialize(&txin.witness).len(),
+                            );
+
+                            let psbt_input_data = if txout.script_pubkey.witness_version().is_some()
+                            {
+                                Box::new(psbt::Input {
                                     witness_utxo: Some(txout.clone()),
                                     non_witness_utxo: Some(prev_tx.as_ref().clone()),
                                     ..Default::default()
-                                }),
-                            },
-                            satisfaction_weight,
+                                })
+                            } else {
+                                Box::new(psbt::Input {
+                                    non_witness_utxo: Some(prev_tx.as_ref().clone()),
+                                    ..Default::default()
+                                })
+                            };
+                            foreign.push(WeightedUtxo {
+                                utxo: Utxo::Foreign {
+                                    outpoint: txin.previous_output,
+                                    sequence: txin.sequence,
+                                    psbt_input: psbt_input_data,
+                                },
+                                satisfaction_weight,
+                            });
+                            true
                         }
-                    }
-                };
+                    };
 
-                Ok(weighted_utxo)
-            })
-            .collect::<Result<Vec<_>, _>>()?;
+                    Ok((local, foreign))
+                },
+            )?;
 
         if tx.output.len() > 1 {
             let mut change_index = None;
@@ -1698,7 +1727,8 @@ impl Wallet {
                 .into_iter()
                 .map(|txout| (txout.script_pubkey, txout.value))
                 .collect(),
-            utxos: original_utxos,
+            utxos: local,
+            foreign_utxos: foreign,
             bumping_fee: Some(tx_builder::PreviousFee {
                 absolute: fee,
                 rate: fee_rate,
@@ -1976,117 +2006,52 @@ impl Wallet {
         descriptor.at_derivation_index(child).ok()
     }
 
-    fn get_available_utxos(&self) -> Vec<(LocalOutput, Weight)> {
-        self.list_unspent()
-            .map(|utxo| {
-                let keychain = utxo.keychain;
-                (utxo, {
-                    self.public_descriptor(keychain)
-                        .max_weight_to_satisfy()
-                        .unwrap()
-                })
-            })
-            .collect()
-    }
-
     /// Given the options returns the list of utxos that must be used to form the
     /// transaction and any further that may be used if needed.
-    fn preselect_utxos(
-        &self,
-        params: &TxParams,
-        current_height: Option<u32>,
-    ) -> (Vec<WeightedUtxo>, Vec<WeightedUtxo>) {
-        let TxParams {
-            change_policy,
-            unspendable,
-            utxos,
-            drain_wallet,
-            manually_selected_only,
-            bumping_fee,
-            ..
-        } = params;
-
-        let manually_selected = utxos.clone();
-        // we mandate confirmed transactions if we're bumping the fee
-        let must_only_use_confirmed_tx = bumping_fee.is_some();
-        let must_use_all_available = *drain_wallet;
-
-        //    must_spend <- manually selected utxos
-        //    may_spend  <- all other available utxos
-        let mut may_spend = self.get_available_utxos();
-
-        may_spend.retain(|may_spend| {
-            !manually_selected
-                .iter()
-                .any(|manually_selected| manually_selected.utxo.outpoint() == may_spend.0.outpoint)
-        });
-        let mut must_spend = manually_selected;
-
-        // NOTE: we are intentionally ignoring `unspendable` here. i.e manual
-        // selection overrides unspendable.
-        if *manually_selected_only {
-            return (must_spend, vec![]);
-        }
-
-        let satisfies_confirmed = may_spend
-            .iter()
-            .map(|u| -> bool {
-                let txid = u.0.outpoint.txid;
-                let tx = match self.indexed_graph.graph().get_tx(txid) {
-                    Some(tx) => tx,
-                    None => return false,
-                };
-
-                // Whether the UTXO is mature and, if needed, confirmed
-                let mut spendable = true;
-                let chain_position = u.0.chain_position;
-                if must_only_use_confirmed_tx && !chain_position.is_confirmed() {
-                    return false;
-                }
-                if tx.is_coinbase() {
-                    debug_assert!(
-                        chain_position.is_confirmed(),
-                        "coinbase must always be confirmed"
-                    );
-                    if let Some(current_height) = current_height {
-                        match chain_position {
-                            ChainPosition::Confirmed { anchor, .. } => {
-                                // https://github.com/bitcoin/bitcoin/blob/c5e67be03bb06a5d7885c55db1f016fbf2333fe3/src/validation.cpp#L373-L375
-                                let spend_height = current_height + 1;
-                                let coin_age_at_spend_height =
-                                    spend_height.saturating_sub(anchor.block_id.height);
-                                spendable &= coin_age_at_spend_height >= COINBASE_MATURITY;
-                            }
-                            ChainPosition::Unconfirmed { .. } => spendable = false,
-                        }
-                    }
-                }
-                spendable
-            })
-            .collect::<Vec<_>>();
-
-        let mut i = 0;
-        may_spend.retain(|u| {
-            let retain = (self.keychains().count() == 1 || change_policy.is_satisfied_by(&u.0))
-                && !unspendable.contains(&u.0.outpoint)
-                && satisfies_confirmed[i];
-            i += 1;
-            retain
-        });
-
-        let mut may_spend = may_spend
-            .into_iter()
-            .map(|(local_utxo, satisfaction_weight)| WeightedUtxo {
-                satisfaction_weight,
-                utxo: Utxo::Local(local_utxo),
-            })
-            .collect();
-
-        if must_use_all_available {
-            must_spend.append(&mut may_spend);
+    fn filter_utxos(&self, params: &TxParams, current_height: u32) -> Vec<WeightedUtxo> {
+        if params.manually_selected_only {
+            vec![]
+        // only process optional UTxOs if manually_selected_only is false
+        } else {
+            self.indexed_graph
+                .graph()
+                // get all unspent UTxOs from wallet
+                // NOTE: the UTxOs returned by the following method already belong to wallet as the
+                // call chain uses get_tx_node infallibly
+                .filter_chain_unspents(
+                    &self.chain,
+                    self.chain.tip().block_id(),
+                    self.indexed_graph.index.outpoints().iter().cloned(),
+                )
+                // only create LocalOutput if UTxO is mature
+                .filter_map(move |((k, i), full_txo)| {
+                    full_txo
+                        .is_mature(current_height)
+                        .then(|| new_local_utxo(k, i, full_txo))
+                })
+                // only process UTxOs not selected manually, they will be considered later in the chain
+                // NOTE: this avoid UTxOs in both required and optional list
+                .filter(|may_spend| !params.utxos.contains(may_spend))
+                // only add to optional UTxOs those which satisfy the change policy if we reuse change
+                .filter(|local_output| {
+                    self.keychains().count() == 1
+                        || params.change_policy.is_satisfied_by(local_output)
+                })
+                // only add to optional UTxOs those marked as spendable
+                .filter(|local_output| !params.unspendable.contains(&local_output.outpoint))
+                // if bumping fees only add to optional UTxOs those confirmed
+                .filter(|local_output| {
+                    params.bumping_fee.is_none() || local_output.chain_position.is_confirmed()
+                })
+                .map(|utxo| WeightedUtxo {
+                    satisfaction_weight: self
+                        .public_descriptor(utxo.keychain)
+                        .max_weight_to_satisfy()
+                        .unwrap(),
+                    utxo: Utxo::Local(utxo),
+                })
+                .collect()
         }
-
-        (must_spend, may_spend)
     }
 
     fn complete_transaction(

crates/wallet/src/wallet/tx_builder.rs

@@ -125,7 +125,8 @@ pub(crate) struct TxParams {
     pub(crate) fee_policy: Option<FeePolicy>,
     pub(crate) internal_policy_path: Option<BTreeMap<String, Vec<usize>>>,
     pub(crate) external_policy_path: Option<BTreeMap<String, Vec<usize>>>,
-    pub(crate) utxos: Vec<WeightedUtxo>,
+    pub(crate) utxos: HashSet<LocalOutput>,
+    pub(crate) foreign_utxos: Vec<WeightedUtxo>,
     pub(crate) unspendable: HashSet<OutPoint>,
     pub(crate) manually_selected_only: bool,
     pub(crate) sighash: Option<psbt::PsbtSighashType>,
@@ -274,26 +275,14 @@ impl<'a, Cs> TxBuilder<'a, Cs> {
     /// These have priority over the "unspendable" utxos, meaning that if a utxo is present both in
     /// the "utxos" and the "unspendable" list, it will be spent.
     pub fn add_utxos(&mut self, outpoints: &[OutPoint]) -> Result<&mut Self, AddUtxoError> {
-        {
-            let wallet = &mut self.wallet;
-            let utxos = outpoints
-                .iter()
-                .map(|outpoint| {
-                    wallet
-                        .get_utxo(*outpoint)
-                        .ok_or(AddUtxoError::UnknownUtxo(*outpoint))
+        let _ = outpoints.iter().try_for_each(|outpoint| {
+            self.wallet
+                .get_utxo(*outpoint)
+                .map(|output| {
+                    self.params.utxos.insert(output);
                 })
-                .collect::<Result<Vec<_>, _>>()?;
-
-            for utxo in utxos {
-                let descriptor = wallet.public_descriptor(utxo.keychain);
-                let satisfaction_weight = descriptor.max_weight_to_satisfy().unwrap();
-                self.params.utxos.push(WeightedUtxo {
-                    satisfaction_weight,
-                    utxo: Utxo::Local(utxo),
-                });
-            }
-        }
+                .ok_or(AddUtxoError::UnknownUtxo(*outpoint))
+        });
 
         Ok(self)
     }
@@ -393,7 +382,7 @@ impl<'a, Cs> TxBuilder<'a, Cs> {
             }
         }
 
-        self.params.utxos.push(WeightedUtxo {
+        self.params.foreign_utxos.push(WeightedUtxo {
             satisfaction_weight,
             utxo: Utxo::Foreign {
                 outpoint,