From a5713a834832b31c7d7a17d19039c7897749ad7d Mon Sep 17 00:00:00 2001 From: Alekos Filini Date: Wed, 14 Oct 2020 15:21:22 +0200 Subject: [PATCH] [wallet] Improve `CoinSelectionAlgorithm` Implement the improvements described in issue #121. Closes #121, closes #131. --- src/database/any.rs | 1 + src/wallet/coin_selection.rs | 155 ++++++++++++++++++++--------------- src/wallet/mod.rs | 83 ++++++++++--------- src/wallet/rbf.rs | 19 +++-- src/wallet/tx_builder.rs | 47 +++++++++-- 5 files changed, 185 insertions(+), 120 deletions(-) diff --git a/src/database/any.rs b/src/database/any.rs index aea44eea..805f8eab 100644 --- a/src/database/any.rs +++ b/src/database/any.rs @@ -87,6 +87,7 @@ macro_rules! impl_inner_method { /// It allows switching database type at runtime. /// /// See [this module](crate::database::any)'s documentation for a usage example. +#[derive(Debug)] pub enum AnyDatabase { Memory(memory::MemoryDatabase), #[cfg(feature = "key-value-db")] diff --git a/src/wallet/coin_selection.rs b/src/wallet/coin_selection.rs index 48893099..7de2b0a2 100644 --- a/src/wallet/coin_selection.rs +++ b/src/wallet/coin_selection.rs @@ -44,37 +44,40 @@ //! # use bitcoin::*; //! # use bitcoin::consensus::serialize; //! # use bdk::wallet::coin_selection::*; +//! # use bdk::database::Database; //! # use bdk::*; //! #[derive(Debug)] //! struct AlwaysSpendEverything; //! -//! impl CoinSelectionAlgorithm for AlwaysSpendEverything { +//! impl CoinSelectionAlgorithm for AlwaysSpendEverything { //! fn coin_select( //! &self, -//! must_use_utxos: Vec, -//! may_use_utxos: Vec, +//! database: &D, +//! must_use_utxos: Vec<(UTXO, usize)>, +//! may_use_utxos: Vec<(UTXO, usize)>, //! fee_rate: FeeRate, //! amount_needed: u64, -//! input_witness_weight: usize, //! fee_amount: f32, //! ) -> Result { //! let mut selected_amount = 0; +//! let mut additional_weight = 0; //! let all_utxos_selected = must_use_utxos //! .into_iter().chain(may_use_utxos) -//! .scan(&mut selected_amount, |selected_amount, utxo| { +//! .scan((&mut selected_amount, &mut additional_weight), |(selected_amount, additional_weight), (utxo, weight)| { +//! let txin = TxIn { +//! previous_output: utxo.outpoint, +//! ..Default::default() +//! }; +//! //! **selected_amount += utxo.txout.value; +//! **additional_weight += serialize(&txin).len() * 4 + weight; +//! //! Some(( -//! TxIn { -//! previous_output: utxo.outpoint, -//! ..Default::default() -//! }, +//! txin, //! utxo.txout.script_pubkey, //! )) //! }) //! .collect::>(); -//! let additional_weight = all_utxos_selected.iter().fold(0, |acc, (txin, _)| { -//! acc + serialize(txin).len() * 4 + input_witness_weight -//! }); //! let additional_fees = additional_weight as f32 * fee_rate.as_sat_vb() / 4.0; //! //! if (fee_amount + additional_fees).ceil() as u64 + amount_needed > selected_amount { @@ -106,6 +109,7 @@ use bitcoin::consensus::encode::serialize; use bitcoin::{Script, TxIn}; +use crate::database::Database; use crate::error::Error; use crate::types::{FeeRate, UTXO}; @@ -130,22 +134,26 @@ pub struct CoinSelectionResult { /// selection algorithm when it creates transactions. /// /// For an example see [this module](crate::wallet::coin_selection)'s documentation. -pub trait CoinSelectionAlgorithm: std::fmt::Debug { +pub trait CoinSelectionAlgorithm: std::fmt::Debug { /// Perform the coin selection /// - /// - `must_use_utxos`: the utxos that must be spent regardless of `amount_needed` - /// - `may_be_spent`: the utxos that may be spent to satisfy `amount_needed` + /// - `database`: a reference to the wallet's database that can be used to lookup additional + /// details for a specific UTXO + /// - `must_use_utxos`: the utxos that must be spent regardless of `amount_needed` with their + /// weight cost + /// - `may_be_spent`: the utxos that may be spent to satisfy `amount_needed` with their weight + /// cost /// - `fee_rate`: fee rate to use /// - `amount_needed`: the amount in satoshi to select - /// - `input_witness_weight`: the weight of an input's witness to keep into account for the fees - /// - `fee_amount`: the amount of fees in satoshi already accumulated from adding outputs + /// - `fee_amount`: the amount of fees in satoshi already accumulated from adding outputs and + /// the transaction's header fn coin_select( &self, - must_use_utxos: Vec, - may_use_utxos: Vec, + database: &D, + must_use_utxos: Vec<(UTXO, usize)>, + may_use_utxos: Vec<(UTXO, usize)>, fee_rate: FeeRate, amount_needed: u64, - input_witness_weight: usize, fee_amount: f32, ) -> Result; } @@ -157,32 +165,33 @@ pub trait CoinSelectionAlgorithm: std::fmt::Debug { #[derive(Debug, Default)] pub struct DumbCoinSelection; -impl CoinSelectionAlgorithm for DumbCoinSelection { +impl CoinSelectionAlgorithm for DumbCoinSelection { fn coin_select( &self, - must_use_utxos: Vec, - mut may_use_utxos: Vec, + _database: &D, + must_use_utxos: Vec<(UTXO, usize)>, + mut may_use_utxos: Vec<(UTXO, usize)>, fee_rate: FeeRate, - outgoing_amount: u64, - input_witness_weight: usize, + amount_needed: u64, mut fee_amount: f32, ) -> Result { let calc_fee_bytes = |wu| (wu as f32) * fee_rate.as_sat_vb() / 4.0; log::debug!( - "outgoing_amount = `{}`, fee_amount = `{}`, fee_rate = `{:?}`", - outgoing_amount, + "amount_needed = `{}`, fee_amount = `{}`, fee_rate = `{:?}`", + amount_needed, fee_amount, fee_rate ); - // We put the "must_use" UTXOs first and make sure the "may_use" are sorted largest to smallest + // We put the "must_use" UTXOs first and make sure the "may_use" are sorted, initially + // smallest to largest, before being reversed with `.rev()`. let utxos = { - may_use_utxos.sort_by(|a, b| b.txout.value.partial_cmp(&a.txout.value).unwrap()); + may_use_utxos.sort_unstable_by_key(|(utxo, _)| utxo.txout.value); must_use_utxos .into_iter() .map(|utxo| (true, utxo)) - .chain(may_use_utxos.into_iter().map(|utxo| (false, utxo))) + .chain(may_use_utxos.into_iter().rev().map(|utxo| (false, utxo))) }; // Keep including inputs until we've got enough. @@ -191,9 +200,8 @@ impl CoinSelectionAlgorithm for DumbCoinSelection { let txin = utxos .scan( (&mut selected_amount, &mut fee_amount), - |(selected_amount, fee_amount), (must_use, utxo)| { - if must_use || **selected_amount < outgoing_amount + (fee_amount.ceil() as u64) - { + |(selected_amount, fee_amount), (must_use, (utxo, weight))| { + if must_use || **selected_amount < amount_needed + (fee_amount.ceil() as u64) { let new_in = TxIn { previous_output: utxo.outpoint, script_sig: Script::default(), @@ -201,8 +209,7 @@ impl CoinSelectionAlgorithm for DumbCoinSelection { witness: vec![], }; - **fee_amount += - calc_fee_bytes(serialize(&new_in).len() * 4 + input_witness_weight); + **fee_amount += calc_fee_bytes(serialize(&new_in).len() * 4 + weight); **selected_amount += utxo.txout.value; log::debug!( @@ -219,7 +226,7 @@ impl CoinSelectionAlgorithm for DumbCoinSelection { ) .collect::>(); - if selected_amount < outgoing_amount + (fee_amount.ceil() as u64) { + if selected_amount < amount_needed + (fee_amount.ceil() as u64) { return Err(Error::InsufficientFunds); } @@ -238,48 +245,56 @@ mod test { use bitcoin::{OutPoint, Script, TxOut}; use super::*; + use crate::database::MemoryDatabase; use crate::types::*; const P2WPKH_WITNESS_SIZE: usize = 73 + 33 + 2; - fn get_test_utxos() -> Vec { + fn get_test_utxos() -> Vec<(UTXO, usize)> { vec![ - UTXO { - outpoint: OutPoint::from_str( - "ebd9813ecebc57ff8f30797de7c205e3c7498ca950ea4341ee51a685ff2fa30a:0", - ) - .unwrap(), - txout: TxOut { - value: 100_000, - script_pubkey: Script::new(), + ( + UTXO { + outpoint: OutPoint::from_str( + "ebd9813ecebc57ff8f30797de7c205e3c7498ca950ea4341ee51a685ff2fa30a:0", + ) + .unwrap(), + txout: TxOut { + value: 100_000, + script_pubkey: Script::new(), + }, + is_internal: false, }, - is_internal: false, - }, - UTXO { - outpoint: OutPoint::from_str( - "65d92ddff6b6dc72c89624a6491997714b90f6004f928d875bc0fd53f264fa85:0", - ) - .unwrap(), - txout: TxOut { - value: 200_000, - script_pubkey: Script::new(), + P2WPKH_WITNESS_SIZE, + ), + ( + UTXO { + outpoint: OutPoint::from_str( + "65d92ddff6b6dc72c89624a6491997714b90f6004f928d875bc0fd53f264fa85:0", + ) + .unwrap(), + txout: TxOut { + value: 200_000, + script_pubkey: Script::new(), + }, + is_internal: true, }, - is_internal: true, - }, + P2WPKH_WITNESS_SIZE, + ), ] } #[test] fn test_dumb_coin_selection_success() { let utxos = get_test_utxos(); + let database = MemoryDatabase::default(); - let result = DumbCoinSelection + let result = DumbCoinSelection::default() .coin_select( - vec![], + &database, utxos, + vec![], FeeRate::from_sat_per_vb(1.0), 250_000, - P2WPKH_WITNESS_SIZE, 50.0, ) .unwrap(); @@ -292,14 +307,15 @@ mod test { #[test] fn test_dumb_coin_selection_use_all() { let utxos = get_test_utxos(); + let database = MemoryDatabase::default(); - let result = DumbCoinSelection + let result = DumbCoinSelection::default() .coin_select( + &database, utxos, vec![], FeeRate::from_sat_per_vb(1.0), 20_000, - P2WPKH_WITNESS_SIZE, 50.0, ) .unwrap(); @@ -312,14 +328,15 @@ mod test { #[test] fn test_dumb_coin_selection_use_only_necessary() { let utxos = get_test_utxos(); + let database = MemoryDatabase::default(); - let result = DumbCoinSelection + let result = DumbCoinSelection::default() .coin_select( + &database, vec![], utxos, FeeRate::from_sat_per_vb(1.0), 20_000, - P2WPKH_WITNESS_SIZE, 50.0, ) .unwrap(); @@ -333,14 +350,15 @@ mod test { #[should_panic(expected = "InsufficientFunds")] fn test_dumb_coin_selection_insufficient_funds() { let utxos = get_test_utxos(); + let database = MemoryDatabase::default(); - DumbCoinSelection + DumbCoinSelection::default() .coin_select( + &database, vec![], utxos, FeeRate::from_sat_per_vb(1.0), 500_000, - P2WPKH_WITNESS_SIZE, 50.0, ) .unwrap(); @@ -350,14 +368,15 @@ mod test { #[should_panic(expected = "InsufficientFunds")] fn test_dumb_coin_selection_insufficient_funds_high_fees() { let utxos = get_test_utxos(); + let database = MemoryDatabase::default(); - DumbCoinSelection + DumbCoinSelection::default() .coin_select( + &database, vec![], utxos, FeeRate::from_sat_per_vb(1000.0), 250_000, - P2WPKH_WITNESS_SIZE, 50.0, ) .unwrap(); diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs index 945b2e18..feb92452 100644 --- a/src/wallet/mod.rs +++ b/src/wallet/mod.rs @@ -241,9 +241,9 @@ where /// // sign and broadcast ... /// # Ok::<(), bdk::Error>(()) /// ``` - pub fn create_tx( + pub fn create_tx>( &self, - builder: TxBuilder, + builder: TxBuilder, ) -> Result<(PSBT, TransactionDetails), Error> { if builder.recipients.is_empty() { return Err(Error::NoAddressees); @@ -334,17 +334,6 @@ where outgoing += value; } - // TODO: use the right weight instead of the maximum, and only fall-back to it if the - // script is unknown in the database - let input_witness_weight = std::cmp::max( - self.get_descriptor_for_script_type(ScriptType::Internal) - .0 - .max_satisfaction_weight(), - self.get_descriptor_for_script_type(ScriptType::External) - .0 - .max_satisfaction_weight(), - ); - if builder.change_policy != tx_builder::ChangeSpendPolicy::ChangeAllowed && self.change_descriptor.is_none() { @@ -369,11 +358,11 @@ where selected_amount, mut fee_amount, } = builder.coin_selection.coin_select( + self.database.borrow().deref(), must_use_utxos, may_use_utxos, fee_rate, outgoing, - input_witness_weight, fee_amount, )?; let (mut txin, prev_script_pubkeys): (Vec<_>, Vec<_>) = txin.into_iter().unzip(); @@ -474,10 +463,10 @@ where // TODO: support for merging multiple transactions while bumping the fees // TODO: option to force addition of an extra output? seems bad for privacy to update the // change - pub fn bump_fee( + pub fn bump_fee>( &self, txid: &Txid, - builder: TxBuilder, + builder: TxBuilder, ) -> Result<(PSBT, TransactionDetails), Error> { let mut details = match self.database.borrow().get_tx(&txid, true)? { None => return Err(Error::TransactionNotFound), @@ -515,7 +504,7 @@ where let mut change_output = None; for (index, txout) in tx.output.iter().enumerate() { // look for an output that we know and that has the right ScriptType. We use - // `get_deget_descriptor_for` to find what's the ScriptType for `Internal` + // `get_descriptor_for` to find what's the ScriptType for `Internal` // addresses really is, because if there's no change_descriptor it's actually equal // to "External" let (_, change_type) = self.get_descriptor_for_script_type(ScriptType::Internal); @@ -533,8 +522,8 @@ where } // we need a change output, add one here and take into account the extra fees for it - if change_output.is_none() { - let change_script = self.get_change_address()?; + let change_script = self.get_change_address()?; + change_output.unwrap_or_else(|| { let change_txout = TxOut { script_pubkey: change_script, value: 0, @@ -543,10 +532,8 @@ where (serialize(&change_txout).len() as f32 * new_feerate.as_sat_vb()).ceil() as u64; tx.output.push(change_txout); - change_output = Some(tx.output.len() - 1); - } - - change_output.unwrap() + tx.output.len() - 1 + }) }; // if `builder.utxos` is Some(_) we have to add inputs and we skip down to the last branch @@ -581,17 +568,6 @@ where let needs_more_inputs = builder.utxos.is_some() || removed_change_output.value <= fee_difference; let added_amount = if needs_more_inputs { - // TODO: use the right weight instead of the maximum, and only fall-back to it if the - // script is unknown in the database - let input_witness_weight = std::cmp::max( - self.get_descriptor_for_script_type(ScriptType::Internal) - .0 - .max_satisfaction_weight(), - self.get_descriptor_for_script_type(ScriptType::External) - .0 - .max_satisfaction_weight(), - ); - let (available_utxos, use_all_utxos) = self.get_available_utxos( builder.change_policy, &builder.utxos, @@ -613,11 +589,11 @@ where selected_amount, fee_amount, } = builder.coin_selection.coin_select( + self.database.borrow().deref(), must_use_utxos, may_use_utxos, new_feerate, fee_difference.saturating_sub(removed_change_output.value), - input_witness_weight, 0.0, )?; fee_difference += fee_amount.ceil() as u64; @@ -945,21 +921,45 @@ where utxo: &Option>, unspendable: &Option>, send_all: bool, - ) -> Result<(Vec, bool), Error> { + ) -> Result<(Vec<(UTXO, usize)>, bool), Error> { let unspendable_set = match unspendable { None => HashSet::new(), Some(vec) => vec.iter().collect(), }; + let external_weight = self + .get_descriptor_for_script_type(ScriptType::External) + .0 + .max_satisfaction_weight(); + let internal_weight = self + .get_descriptor_for_script_type(ScriptType::Internal) + .0 + .max_satisfaction_weight(); + + let add_weight = |utxo: UTXO| { + let weight = match utxo.is_internal { + true => internal_weight, + false => external_weight, + }; + + (utxo, weight) + }; + match utxo { // with manual coin selection we always want to spend all the selected utxos, no matter // what (even if they are marked as unspendable) Some(raw_utxos) => { let full_utxos = raw_utxos .iter() - .map(|u| self.database.borrow().get_utxo(&u)) - .collect::>, _>>()? - .ok_or(Error::UnknownUTXO)?; + .map(|u| { + Ok(add_weight( + self.database + .borrow() + .get_utxo(&u)? + .ok_or(Error::UnknownUTXO)?, + )) + }) + .collect::, Error>>()?; Ok((full_utxos, true)) } @@ -971,6 +971,7 @@ where Ok(( utxos .filter(|u| !unspendable_set.contains(&u.outpoint)) + .map(add_weight) .collect(), send_all, )) @@ -978,11 +979,11 @@ where } } - fn complete_transaction( + fn complete_transaction>( &self, tx: Transaction, prev_script_pubkeys: HashMap, - builder: TxBuilder, + builder: TxBuilder, ) -> Result { let mut psbt = PSBT::from_unsigned_tx(tx)?; diff --git a/src/wallet/rbf.rs b/src/wallet/rbf.rs index 42a7b00b..6080641e 100644 --- a/src/wallet/rbf.rs +++ b/src/wallet/rbf.rs @@ -27,16 +27,16 @@ use crate::error::Error; use crate::types::*; /// Filters unspent utxos -pub(super) fn filter_available, D: Database>( +pub(super) fn filter_available, D: Database>( database: &D, iter: I, -) -> Result, Error> { +) -> Result, Error> { Ok(iter - .map(|utxo| { + .map(|(utxo, weight)| { Ok(match database.get_tx(&utxo.outpoint.txid, true)? { None => None, Some(tx) if tx.height.is_none() => None, - Some(_) => Some(utxo), + Some(_) => Some((utxo, weight)), }) }) .collect::, Error>>()? @@ -120,8 +120,15 @@ mod test { vec![50_000], ); - let filtered = - filter_available(&database, database.iter_utxos().unwrap().into_iter()).unwrap(); + let filtered = filter_available( + &database, + database + .iter_utxos() + .unwrap() + .into_iter() + .map(|utxo| (utxo, 0)), + ) + .unwrap(); assert_eq!(filtered, &[]); } } diff --git a/src/wallet/tx_builder.rs b/src/wallet/tx_builder.rs index 017d26f7..309825bc 100644 --- a/src/wallet/tx_builder.rs +++ b/src/wallet/tx_builder.rs @@ -38,14 +38,17 @@ //! .fee_rate(FeeRate::from_sat_per_vb(5.0)) //! .do_not_spend_change() //! .enable_rbf(); +//! # let builder: TxBuilder = builder; //! ``` use std::collections::BTreeMap; use std::default::Default; +use std::marker::PhantomData; use bitcoin::{OutPoint, Script, SigHashType, Transaction}; use super::coin_selection::{CoinSelectionAlgorithm, DefaultCoinSelectionAlgorithm}; +use crate::database::Database; use crate::types::{FeeRate, UTXO}; /// A transaction builder @@ -53,8 +56,8 @@ use crate::types::{FeeRate, UTXO}; /// This structure contains the configuration that the wallet must follow to build a transaction. /// /// For an example see [this module](super::tx_builder)'s documentation; -#[derive(Debug, Default)] -pub struct TxBuilder { +#[derive(Debug)] +pub struct TxBuilder> { pub(crate) recipients: Vec<(Script, u64)>, pub(crate) send_all: bool, pub(crate) fee_rate: Option, @@ -69,9 +72,38 @@ pub struct TxBuilder { pub(crate) change_policy: ChangeSpendPolicy, pub(crate) force_non_witness_utxo: bool, pub(crate) coin_selection: Cs, + + phantom: PhantomData, } -impl TxBuilder { +// Unfortunately derive doesn't work with `PhantomData`: https://github.com/rust-lang/rust/issues/26925 +impl> Default for TxBuilder +where + Cs: Default, +{ + fn default() -> Self { + TxBuilder { + recipients: Default::default(), + send_all: Default::default(), + fee_rate: Default::default(), + policy_path: Default::default(), + utxos: Default::default(), + unspendable: Default::default(), + sighash: Default::default(), + ordering: Default::default(), + locktime: Default::default(), + rbf: Default::default(), + version: Default::default(), + change_policy: Default::default(), + force_non_witness_utxo: Default::default(), + coin_selection: Default::default(), + + phantom: PhantomData, + } + } +} + +impl TxBuilder { /// Create an empty builder pub fn new() -> Self { Self::default() @@ -83,7 +115,7 @@ impl TxBuilder { } } -impl TxBuilder { +impl> TxBuilder { /// Replace the recipients already added with a new list pub fn set_recipients(mut self, recipients: Vec<(Script, u64)>) -> Self { self.recipients = recipients; @@ -248,7 +280,10 @@ impl TxBuilder { /// Choose the coin selection algorithm /// /// Overrides the [`DefaultCoinSelectionAlgorithm`](super::coin_selection::DefaultCoinSelectionAlgorithm). - pub fn coin_selection(self, coin_selection: P) -> TxBuilder

{ + pub fn coin_selection>( + self, + coin_selection: P, + ) -> TxBuilder { TxBuilder { recipients: self.recipients, send_all: self.send_all, @@ -264,6 +299,8 @@ impl TxBuilder { change_policy: self.change_policy, force_non_witness_utxo: self.force_non_witness_utxo, coin_selection, + + phantom: PhantomData, } } }