Skip to content

Commit

Permalink
Return Iterator instead of Vector for SRD
Browse files Browse the repository at this point in the history
  • Loading branch information
yancyribbens committed Feb 13, 2024
1 parent 7ebc5bc commit 4f798b8
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 37 deletions.
10 changes: 3 additions & 7 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,15 +59,11 @@ pub struct WeightedUtxo {
#[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
pub fn select_coins<T: Utxo>(
target: Amount,
cost_of_change: u64,
_cost_of_change: u64,
fee_rate: FeeRate,
weighted_utxos: &mut [WeightedUtxo],
utxo_pool: &mut [T],
) -> Option<Vec<WeightedUtxo>> {
match select_coins_bnb(target.to_sat(), cost_of_change, utxo_pool) {
Some(_res) => Some(Vec::new()),
None => select_coins_srd(target, fee_rate, weighted_utxos, &mut thread_rng()),
}
) -> Option<impl Iterator<Item = &'_ WeightedUtxo>> {
select_coins_srd(target, fee_rate, weighted_utxos, &mut thread_rng())
}

/// Select coins using BnB algorithm similar to what is done in bitcoin
Expand Down
80 changes: 50 additions & 30 deletions src/single_random_draw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,20 +28,22 @@ use rand::seq::SliceRandom;
/// /// * `fee_rate` - ratio of transaction amount per size.
/// /// * `weighted_utxos` - Weighted UTXOs from which to sum the target amount.
/// /// * `rng` - used primarily by tests to make the selection deterministic.
pub fn select_coins_srd<R: rand::Rng + ?Sized>(
pub fn select_coins_srd<'a, R: rand::Rng + ?Sized>(
target: Amount,
fee_rate: FeeRate,
weighted_utxos: &mut [WeightedUtxo],
weighted_utxos: &'a mut [WeightedUtxo],
rng: &mut R,
) -> Option<Vec<WeightedUtxo>> {
let mut result: Vec<WeightedUtxo> = Vec::new();
) -> Option<impl Iterator<Item = &'a WeightedUtxo>> {
let mut result: Vec<_> = weighted_utxos.iter().collect();
let mut origin = result.to_owned();
origin.shuffle(rng);

weighted_utxos.shuffle(rng);
result.clear();

let threshold = target + CHANGE_LOWER;
let mut value = Amount::ZERO;

for w_utxo in weighted_utxos {
for w_utxo in origin {
let utxo_value = w_utxo.utxo.value;
let effective_value = effective_value(fee_rate, w_utxo.satisfaction_weight, utxo_value)?;

Expand All @@ -50,14 +52,14 @@ pub fn select_coins_srd<R: rand::Rng + ?Sized>(
Err(_) => continue,
};

result.push(w_utxo.clone());
result.push(w_utxo);

if value >= threshold {
return Some(result);
return Some(result.into_iter());
}
}

Some(Vec::new())
None
}

#[cfg(test)]
Expand Down Expand Up @@ -115,32 +117,41 @@ mod tests {
let target: Amount = Amount::from_str("1.5 cBTC").unwrap();
let mut weighted_utxos: Vec<WeightedUtxo> = create_weighted_utxos();

let result = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error");
let result: Vec<&WeightedUtxo> =
select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error")
.collect();

assert_eq!(vec![weighted_utxos[0].clone()], result);
let expected_result = Amount::from_str("2 cBTC").unwrap();
assert_eq!(result.len(), 1);
assert_eq!(expected_result, result[0].utxo.value);
}

#[test]
fn select_coins_srd_no_solution() {
let target: Amount = Amount::from_str("4 cBTC").unwrap();
let mut weighted_utxos: Vec<WeightedUtxo> = create_weighted_utxos();

let result = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error");

assert!(result.is_empty());
let result = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut get_rng());
assert!(result.is_none())
}

#[test]
fn select_coins_srd_all_solution() {
let target: Amount = Amount::from_str("2.5 cBTC").unwrap();
let mut weighted_utxos: Vec<WeightedUtxo> = create_weighted_utxos();

let result = select_coins_srd(target, FeeRate::ZERO, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error");
let result: Vec<&WeightedUtxo> =
select_coins_srd(target, FeeRate::ZERO, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error")
.collect();

let expected_second_element = Amount::from_str("1 cBTC").unwrap();
let expected_first_element = Amount::from_str("2 cBTC").unwrap();

assert_eq!(weighted_utxos.clone(), result);
assert_eq!(result.len(), 2);
assert_eq!(result[0].utxo.value, expected_first_element);
assert_eq!(result[1].utxo.value, expected_second_element);
}

#[test]
Expand All @@ -157,12 +168,16 @@ mod tests {
});

let mut rng = get_rng();
let result = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut rng)
.expect("unexpected error");
let result: Vec<_> = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut rng)
.expect("unexpected error")
.collect();

let mut expected_utxos = create_weighted_utxos();
expected_utxos.shuffle(&mut rng);
assert_eq!(result, expected_utxos);
let expected_second_element = Amount::from_str("1 cBTC").unwrap();
let expected_first_element = Amount::from_str("2 cBTC").unwrap();

assert_eq!(result.len(), 2);
assert_eq!(result[0].utxo.value, expected_first_element);
assert_eq!(result[1].utxo.value, expected_second_element);
}

#[test]
Expand All @@ -179,10 +194,9 @@ mod tests {
let target: Amount = Amount::from_str("3 cBTC").unwrap();
let mut weighted_utxos: Vec<WeightedUtxo> = create_weighted_utxos();

let result = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error");
let result = select_coins_srd(target, FEE_RATE, &mut weighted_utxos, &mut get_rng());

assert!(result.is_empty());
assert!(result.is_none());
}

#[test]
Expand All @@ -197,10 +211,16 @@ mod tests {
let fee_rate: FeeRate = FeeRate::from_sat_per_kwu(40);
let mut weighted_utxos: Vec<WeightedUtxo> = create_weighted_utxos();

let result = select_coins_srd(target, fee_rate, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error");
let result: Vec<_> =
select_coins_srd(target, fee_rate, &mut weighted_utxos, &mut get_rng())
.expect("unexpected error")
.collect();
let expected_second_element = Amount::from_str("1 cBTC").unwrap();
let expected_first_element = Amount::from_str("2 cBTC").unwrap();

assert_eq!(weighted_utxos.clone(), result);
assert_eq!(result.len(), 2);
assert_eq!(result[0].utxo.value, expected_first_element);
assert_eq!(result[1].utxo.value, expected_second_element);
}

#[test]
Expand Down

0 comments on commit 4f798b8

Please sign in to comment.