Skip to content

Commit

Permalink
test(coin_selection): add test for deterministic utxo selection
Browse files Browse the repository at this point in the history 
Added back ignored branch and bounnd tests and cleaned up calculation for target amounts.
  • Loading branch information
@notmandatory
notmandatory committed Sep 11, 2024
1 parent c18204d commit 65be4ea
Showing 1 changed file with 136 additions and 30 deletions.
166 changes: 136 additions & 30 deletions crates/wallet/src/wallet/coin_selection.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,6 @@ pub trait CoinSelectionAlgorithm: core::fmt::Debug {
/// accumulated from added outputs and transaction’s header.
/// - `drain_script`: the script to use in case of change
/// - `rand`: random number generated used by some coin selection algorithms such as [`SingleRandomDraw`]
#[allow(clippy::too_many_arguments)]
fn coin_select<R: RngCore>(
&self,
required_utxos: Vec<WeightedUtxo>,
Expand Down Expand Up @@ -398,14 +397,14 @@ impl OutputGroup {
///
/// Code adapted from Bitcoin Core's implementation and from Mark Erhardt Master's Thesis: <http://murch.one/wp-content/uploads/2016/11/erhardt2016coinselection.pdf>
#[derive(Debug, Clone)]
pub struct BranchAndBoundCoinSelection<FA = SingleRandomDraw> {
pub struct BranchAndBoundCoinSelection<Cs = SingleRandomDraw> {
size_of_change: u64,
fallback_algorithm: FA,
fallback_algorithm: Cs,
}

/// Error returned by branch and bond coin selection.
#[derive(Debug)]
enum BnBError {
enum BnbError {
/// Branch and bound coin selection tries to avoid needing a change by finding the right inputs for
/// the desired outputs plus fee, if there is not such combination this error is thrown
NoExactMatch,
Expand All @@ -414,19 +413,19 @@ enum BnBError {
TotalTriesExceeded,
}

impl<FA: Default> Default for BranchAndBoundCoinSelection<FA> {
impl<Cs: Default> Default for BranchAndBoundCoinSelection<Cs> {
fn default() -> Self {
Self {
// P2WPKH cost of change -> value (8 bytes) + script len (1 bytes) + script (22 bytes)
size_of_change: 8 + 1 + 22,
fallback_algorithm: FA::default(),
fallback_algorithm: Cs::default(),
}
}
}

impl<FA> BranchAndBoundCoinSelection<FA> {
impl<Cs> BranchAndBoundCoinSelection<Cs> {
/// Create new instance with a target `size_of_change` and `fallback_algorithm`.
pub fn new(size_of_change: u64, fallback_algorithm: FA) -> Self {
pub fn new(size_of_change: u64, fallback_algorithm: Cs) -> Self {
Self {
size_of_change,
fallback_algorithm,
Expand All @@ -436,7 +435,7 @@ impl<FA> BranchAndBoundCoinSelection<FA> {

const BNB_TOTAL_TRIES: usize = 100_000;

impl<FA: CoinSelectionAlgorithm> CoinSelectionAlgorithm for BranchAndBoundCoinSelection<FA> {
impl<Cs: CoinSelectionAlgorithm> CoinSelectionAlgorithm for BranchAndBoundCoinSelection<Cs> {
fn coin_select<R: RngCore>(
&self,
required_utxos: Vec<WeightedUtxo>,
Expand Down Expand Up @@ -544,7 +543,7 @@ impl<FA: CoinSelectionAlgorithm> CoinSelectionAlgorithm for BranchAndBoundCoinSe
}
}

impl<FA> BranchAndBoundCoinSelection<FA> {
impl<Cs> BranchAndBoundCoinSelection<Cs> {
// TODO: make this more Rust-onic :)
// (And perhaps refactor with less arguments?)
#[allow(clippy::too_many_arguments)]
Expand All @@ -558,7 +557,7 @@ impl<FA> BranchAndBoundCoinSelection<FA> {
cost_of_change: u64,
drain_script: &Script,
fee_rate: FeeRate,
) -> Result<CoinSelectionResult, BnBError> {
) -> Result<CoinSelectionResult, BnbError> {
// current_selection[i] will contain true if we are using optional_utxos[i],
// false otherwise. Note that current_selection.len() could be less than
// optional_utxos.len(), it just means that we still haven't decided if we should keep
Expand Down Expand Up @@ -614,7 +613,7 @@ impl<FA> BranchAndBoundCoinSelection<FA> {
// We have walked back to the first utxo and no branch is untraversed. All solutions searched
// If best selection is empty, then there's no exact match
if best_selection.is_empty() {
return Err(BnBError::NoExactMatch);
return Err(BnbError::NoExactMatch);
}
break;
}
Expand All @@ -641,7 +640,7 @@ impl<FA> BranchAndBoundCoinSelection<FA> {

// Check for solution
if best_selection.is_empty() {
return Err(BnBError::TotalTriesExceeded);
return Err(BnbError::TotalTriesExceeded);
}

// Set output set
Expand Down Expand Up @@ -929,6 +928,14 @@ mod test {
.sum()
}

fn calc_target_amount(utxos: &[WeightedUtxo], fee_rate: FeeRate) -> u64 {
utxos
.iter()
.cloned()
.map(|utxo| u64::try_from(OutputGroup::new(utxo, fee_rate).effective_value).unwrap())
.sum()
}

#[test]
fn test_largest_first_coin_selection_success() {
let utxos = get_test_utxos();
Expand Down Expand Up @@ -1143,7 +1150,6 @@ mod test {
}

#[test]
#[ignore = "SRD fn was moved out of BnB"]
fn test_bnb_coin_selection_success() {
// In this case bnb won't find a suitable match and single random draw will
// select three outputs
Expand Down Expand Up @@ -1190,17 +1196,18 @@ mod test {
}

#[test]
#[ignore = "no exact match for bnb, previously fell back to SRD"]
fn test_bnb_coin_selection_optional_are_enough() {
let utxos = get_test_utxos();
let drain_script = ScriptBuf::default();
let target_amount = 299756 + FEE_AMOUNT;
let fee_rate = FeeRate::BROADCAST_MIN;
// first and third utxo's effective value
let target_amount = calc_target_amount(&[utxos[0].clone(), utxos[2].clone()], fee_rate);

let result = BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.coin_select(
vec![],
utxos,
FeeRate::from_sat_per_vb_unchecked(1),
fee_rate,
target_amount,
&drain_script,
&mut thread_rng(),
Expand Down Expand Up @@ -1233,7 +1240,6 @@ mod test {
}

#[test]
#[ignore]
fn test_bnb_coin_selection_required_not_enough() {
let utxos = get_test_utxos();

Expand All @@ -1252,13 +1258,15 @@ mod test {
assert!(amount > 150_000);
let drain_script = ScriptBuf::default();

let target_amount = 150_000 + FEE_AMOUNT;
let fee_rate = FeeRate::BROADCAST_MIN;
// first and third utxo's effective value
let target_amount = calc_target_amount(&[utxos[0].clone(), utxos[2].clone()], fee_rate);

let result = BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.coin_select(
required,
optional,
FeeRate::from_sat_per_vb_unchecked(1),
fee_rate,
target_amount,
&drain_script,
&mut thread_rng(),
Expand Down Expand Up @@ -1312,14 +1320,15 @@ mod test {
fn test_bnb_coin_selection_check_fee_rate() {
let utxos = get_test_utxos();
let drain_script = ScriptBuf::default();
let target_amount = 99932; // first utxo's effective value
let feerate = FeeRate::BROADCAST_MIN;
let fee_rate = FeeRate::BROADCAST_MIN;
// first utxo's effective value
let target_amount = calc_target_amount(&utxos[0..1], fee_rate);

let result = BranchAndBoundCoinSelection::new(0, SingleRandomDraw)
let result = BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.coin_select(
vec![],
utxos,
feerate,
fee_rate,
target_amount,
&drain_script,
&mut thread_rng(),
Expand All @@ -1332,7 +1341,7 @@ mod test {
TxIn::default().segwit_weight().to_wu() + P2WPKH_SATISFACTION_SIZE as u64;
// the final fee rate should be exactly the same as the fee rate given
let result_feerate = Amount::from_sat(result.fee_amount) / Weight::from_wu(input_weight);
assert_eq!(result_feerate, feerate);
assert_eq!(result_feerate, fee_rate);
}

#[test]
Expand All @@ -1344,7 +1353,7 @@ mod test {
let mut optional_utxos = generate_random_utxos(&mut rng, 16);
let target_amount = sum_random_utxos(&mut rng, &mut optional_utxos);
let drain_script = ScriptBuf::default();
let result = BranchAndBoundCoinSelection::new(0, SingleRandomDraw)
let result = BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.coin_select(
vec![],
optional_utxos,
Expand Down Expand Up @@ -1374,7 +1383,7 @@ mod test {

let drain_script = ScriptBuf::default();
let target_amount = 20_000 + FEE_AMOUNT;
BranchAndBoundCoinSelection::new(size_of_change, SingleRandomDraw)
BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.bnb(
vec![],
utxos,
Expand Down Expand Up @@ -1405,7 +1414,7 @@ mod test {

let drain_script = ScriptBuf::default();

BranchAndBoundCoinSelection::new(size_of_change, SingleRandomDraw)
BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.bnb(
vec![],
utxos,
Expand Down Expand Up @@ -1441,7 +1450,7 @@ mod test {

let drain_script = ScriptBuf::default();

let result = BranchAndBoundCoinSelection::new(size_of_change, SingleRandomDraw)
let result = BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.bnb(
vec![],
utxos,
Expand Down Expand Up @@ -1481,7 +1490,7 @@ mod test {

let drain_script = ScriptBuf::default();

let result = BranchAndBoundCoinSelection::new(0, SingleRandomDraw)
let result = BranchAndBoundCoinSelection::<SingleRandomDraw>::default()
.bnb(
vec![],
optional_utxos,
Expand Down Expand Up @@ -1681,4 +1690,101 @@ mod test {
);
}
}

#[test]
fn test_deterministic_coin_selection_picks_same_utxos() {
enum CoinSelectionAlgo {
BranchAndBound,
OldestFirst,
LargestFirst,
}

struct TestCase<'a> {
name: &'a str,
coin_selection_algo: CoinSelectionAlgo,
exp_vouts: &'a [u32],
}

let test_cases = [
TestCase {
name: "branch and bound",
coin_selection_algo: CoinSelectionAlgo::BranchAndBound,
// note: we expect these to be sorted largest first, which indicates
// BnB succeeded with no fallback
exp_vouts: &[29, 28, 27],
},
TestCase {
name: "oldest first",
coin_selection_algo: CoinSelectionAlgo::OldestFirst,
exp_vouts: &[0, 1, 2],
},
TestCase {
name: "largest first",
coin_selection_algo: CoinSelectionAlgo::LargestFirst,
exp_vouts: &[29, 28, 27],
},
];

let optional = generate_same_value_utxos(100_000, 30);
let fee_rate = FeeRate::from_sat_per_vb_unchecked(1);
let target_amount = calc_target_amount(&optional[0..3], fee_rate);
assert_eq!(target_amount, 299_796);
let drain_script = ScriptBuf::default();

for tc in test_cases {
let optional = optional.clone();

let result = match tc.coin_selection_algo {
CoinSelectionAlgo::BranchAndBound => {
BranchAndBoundCoinSelection::<SingleRandomDraw>::default().coin_select(
vec![],
optional,
fee_rate,
target_amount,
&drain_script,
&mut thread_rng(),
)
}
CoinSelectionAlgo::OldestFirst => OldestFirstCoinSelection.coin_select(
vec![],
optional,
fee_rate,
target_amount,
&drain_script,
&mut thread_rng(),
),
CoinSelectionAlgo::LargestFirst => LargestFirstCoinSelection.coin_select(
vec![],
optional,
fee_rate,
target_amount,
&drain_script,
&mut thread_rng(),
),
};

assert!(result.is_ok(), "coin_select failed {}", tc.name);
let result = result.unwrap();
assert!(matches!(result.excess, Excess::NoChange { .. },));
assert_eq!(
result.selected.len(),
3,
"wrong selected len for {}",
tc.name
);
assert_eq!(
result.selected_amount(),
300_000,
"wrong selected amount for {}",
tc.name
);
assert_eq!(result.fee_amount, 204, "wrong fee amount for {}", tc.name);
let vouts = result
.selected
.iter()
.map(|utxo| utxo.outpoint().vout)
.collect::<Vec<u32>>();
assert_eq!(vouts, tc.exp_vouts, "wrong selected vouts for {}", tc.name);
}
}
}

0 comments on commit 65be4ea

Please sign in to comment.