From 438cd4682d14b5f4c6c9c653c9f05ccaaecb815d Mon Sep 17 00:00:00 2001 From: Leonardo Lima Date: Fri, 7 Jun 2024 16:33:02 -0300 Subject: [PATCH] refactor(wallet)!: change WeightedUtxo to use Weight type --- crates/wallet/src/types.rs | 4 +- crates/wallet/src/wallet/coin_selection.rs | 32 +++++++-------- crates/wallet/src/wallet/mod.rs | 16 ++++---- crates/wallet/src/wallet/tx_builder.rs | 12 +++--- crates/wallet/tests/wallet.rs | 47 +++++----------------- 5 files changed, 43 insertions(+), 68 deletions(-) diff --git a/crates/wallet/src/types.rs b/crates/wallet/src/types.rs index 4ce961b7e..0370e55d6 100644 --- a/crates/wallet/src/types.rs +++ b/crates/wallet/src/types.rs @@ -14,7 +14,7 @@ use core::convert::AsRef; use bdk_chain::ConfirmationTime; use bitcoin::blockdata::transaction::{OutPoint, Sequence, TxOut}; -use bitcoin::psbt; +use bitcoin::{psbt, Weight}; use serde::{Deserialize, Serialize}; @@ -72,7 +72,7 @@ pub struct WeightedUtxo { /// properly maintain the feerate when adding this input to a transaction during coin selection. /// /// [weight units]: https://en.bitcoin.it/wiki/Weight_units - pub satisfaction_weight: usize, + pub satisfaction_weight: Weight, /// The UTXO pub utxo: Utxo, } diff --git a/crates/wallet/src/wallet/coin_selection.rs b/crates/wallet/src/wallet/coin_selection.rs index 4016e05ca..301e57755 100644 --- a/crates/wallet/src/wallet/coin_selection.rs +++ b/crates/wallet/src/wallet/coin_selection.rs @@ -52,11 +52,10 @@ //! (&mut selected_amount, &mut additional_weight), //! |(selected_amount, additional_weight), weighted_utxo| { //! **selected_amount += weighted_utxo.utxo.txout().value.to_sat(); -//! **additional_weight += Weight::from_wu( -//! (TxIn::default().segwit_weight().to_wu() -//! + weighted_utxo.satisfaction_weight as u64) -//! as u64, -//! ); +//! **additional_weight += TxIn::default() +//! .segwit_weight() +//! .checked_add(weighted_utxo.satisfaction_weight) +//! .expect("`Weight` addition should not cause an integer overflow"); //! Some(weighted_utxo.utxo) //! }, //! ) @@ -344,10 +343,10 @@ fn select_sorted_utxos( |(selected_amount, fee_amount), (must_use, weighted_utxo)| { if must_use || **selected_amount < target_amount + **fee_amount { **fee_amount += (fee_rate - * Weight::from_wu( - TxIn::default().segwit_weight().to_wu() - + weighted_utxo.satisfaction_weight as u64, - )) + * (TxIn::default() + .segwit_weight() + .checked_add(weighted_utxo.satisfaction_weight) + .expect("`Weight` addition should not cause an integer overflow"))) .to_sat(); **selected_amount += weighted_utxo.utxo.txout().value.to_sat(); Some(weighted_utxo.utxo) @@ -390,9 +389,10 @@ struct OutputGroup { impl OutputGroup { fn new(weighted_utxo: WeightedUtxo, fee_rate: FeeRate) -> Self { let fee = (fee_rate - * Weight::from_wu( - TxIn::default().segwit_weight().to_wu() + weighted_utxo.satisfaction_weight as u64, - )) + * (TxIn::default() + .segwit_weight() + .checked_add(weighted_utxo.satisfaction_weight) + .expect("`Weight` addition should not cause an integer overflow"))) .to_sat(); let effective_value = weighted_utxo.utxo.txout().value.to_sat() as i64 - fee as i64; OutputGroup { @@ -774,7 +774,7 @@ mod test { )) .unwrap(); WeightedUtxo { - satisfaction_weight: P2WPKH_SATISFACTION_SIZE, + satisfaction_weight: Weight::from_wu_usize(P2WPKH_SATISFACTION_SIZE), utxo: Utxo::Local(LocalOutput { outpoint, txout: TxOut { @@ -834,7 +834,7 @@ mod test { let mut res = Vec::new(); for i in 0..utxos_number { res.push(WeightedUtxo { - satisfaction_weight: P2WPKH_SATISFACTION_SIZE, + satisfaction_weight: Weight::from_wu_usize(P2WPKH_SATISFACTION_SIZE), utxo: Utxo::Local(LocalOutput { outpoint: OutPoint::from_str(&format!( "ebd9813ecebc57ff8f30797de7c205e3c7498ca950ea4341ee51a685ff2fa30a:{}", @@ -865,7 +865,7 @@ mod test { fn generate_same_value_utxos(utxos_value: u64, utxos_number: usize) -> Vec { (0..utxos_number) .map(|i| WeightedUtxo { - satisfaction_weight: P2WPKH_SATISFACTION_SIZE, + satisfaction_weight: Weight::from_wu_usize(P2WPKH_SATISFACTION_SIZE), utxo: Utxo::Local(LocalOutput { outpoint: OutPoint::from_str(&format!( "ebd9813ecebc57ff8f30797de7c205e3c7498ca950ea4341ee51a685ff2fa30a:{}", @@ -1512,7 +1512,7 @@ mod test { fn test_filter_duplicates() { fn utxo(txid: &str, value: u64) -> WeightedUtxo { WeightedUtxo { - satisfaction_weight: 0, + satisfaction_weight: Weight::ZERO, utxo: Utxo::Local(LocalOutput { outpoint: OutPoint::new(bitcoin::hashes::Hash::hash(txid.as_bytes()), 0), txout: TxOut { diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs index 400ecde1d..71fb39f7d 100644 --- a/crates/wallet/src/wallet/mod.rs +++ b/crates/wallet/src/wallet/mod.rs @@ -31,7 +31,6 @@ use bdk_chain::{ Append, BlockId, ChainPosition, ConfirmationTime, ConfirmationTimeHeightAnchor, FullTxOut, Indexed, IndexedTxGraph, }; -use bitcoin::secp256k1::{All, Secp256k1}; use bitcoin::sighash::{EcdsaSighashType, TapSighashType}; use bitcoin::{ absolute, psbt, Address, Block, FeeRate, Network, OutPoint, Script, ScriptBuf, Sequence, @@ -39,6 +38,10 @@ use bitcoin::{ }; use bitcoin::{consensus::encode::serialize, transaction, BlockHash, Psbt}; use bitcoin::{constants::genesis_block, Amount}; +use bitcoin::{ + secp256k1::{All, Secp256k1}, + Weight, +}; use core::fmt; use core::mem; use core::ops::Deref; @@ -1662,8 +1665,7 @@ impl Wallet { let satisfaction_weight = self .get_descriptor_for_keychain(keychain) .max_weight_to_satisfy() - .unwrap() - .to_wu() as usize; + .unwrap(); WeightedUtxo { utxo: Utxo::Local(LocalOutput { outpoint: txin.previous_output, @@ -1677,8 +1679,9 @@ impl Wallet { } } None => { - let satisfaction_weight = - serialize(&txin.script_sig).len() * 4 + serialize(&txin.witness).len(); + 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, @@ -1987,7 +1990,7 @@ impl Wallet { descriptor.at_derivation_index(child).ok() } - fn get_available_utxos(&self) -> Vec<(LocalOutput, usize)> { + fn get_available_utxos(&self) -> Vec<(LocalOutput, Weight)> { self.list_unspent() .map(|utxo| { let keychain = utxo.keychain; @@ -1995,7 +1998,6 @@ impl Wallet { self.get_descriptor_for_keychain(keychain) .max_weight_to_satisfy() .unwrap() - .to_wu() as usize }) }) .collect() diff --git a/crates/wallet/src/wallet/tx_builder.rs b/crates/wallet/src/wallet/tx_builder.rs index 0026025de..e8d6d3d0c 100644 --- a/crates/wallet/src/wallet/tx_builder.rs +++ b/crates/wallet/src/wallet/tx_builder.rs @@ -44,7 +44,9 @@ use core::fmt; use bitcoin::psbt::{self, Psbt}; use bitcoin::script::PushBytes; -use bitcoin::{absolute, Amount, FeeRate, OutPoint, ScriptBuf, Sequence, Transaction, Txid}; +use bitcoin::{ + absolute, Amount, FeeRate, OutPoint, ScriptBuf, Sequence, Transaction, Txid, Weight, +}; use rand_core::RngCore; use super::coin_selection::CoinSelectionAlgorithm; @@ -295,9 +297,7 @@ impl<'a, Cs> TxBuilder<'a, Cs> { for utxo in utxos { let descriptor = wallet.get_descriptor_for_keychain(utxo.keychain); - - let satisfaction_weight = - descriptor.max_weight_to_satisfy().unwrap().to_wu() as usize; + let satisfaction_weight = descriptor.max_weight_to_satisfy().unwrap(); self.params.utxos.push(WeightedUtxo { satisfaction_weight, utxo: Utxo::Local(utxo), @@ -366,7 +366,7 @@ impl<'a, Cs> TxBuilder<'a, Cs> { &mut self, outpoint: OutPoint, psbt_input: psbt::Input, - satisfaction_weight: usize, + satisfaction_weight: Weight, ) -> Result<&mut Self, AddForeignUtxoError> { self.add_foreign_utxo_with_sequence( outpoint, @@ -381,7 +381,7 @@ impl<'a, Cs> TxBuilder<'a, Cs> { &mut self, outpoint: OutPoint, psbt_input: psbt::Input, - satisfaction_weight: usize, + satisfaction_weight: Weight, sequence: Sequence, ) -> Result<&mut Self, AddForeignUtxoError> { if psbt_input.witness_utxo.is_none() { diff --git a/crates/wallet/tests/wallet.rs b/crates/wallet/tests/wallet.rs index ef5575acb..5b8d3abd0 100644 --- a/crates/wallet/tests/wallet.rs +++ b/crates/wallet/tests/wallet.rs @@ -1387,11 +1387,7 @@ fn test_add_foreign_utxo() { builder .add_recipient(addr.script_pubkey(), Amount::from_sat(60_000)) .only_witness_utxo() - .add_foreign_utxo( - utxo.outpoint, - psbt_input, - foreign_utxo_satisfaction.to_wu() as usize, - ) + .add_foreign_utxo(utxo.outpoint, psbt_input, foreign_utxo_satisfaction) .unwrap(); let mut psbt = builder.finish().unwrap(); wallet1.insert_txout(utxo.outpoint, utxo.txout); @@ -1467,11 +1463,7 @@ fn test_calculate_fee_with_missing_foreign_utxo() { builder .add_recipient(addr.script_pubkey(), Amount::from_sat(60_000)) .only_witness_utxo() - .add_foreign_utxo( - utxo.outpoint, - psbt_input, - foreign_utxo_satisfaction.to_wu() as usize, - ) + .add_foreign_utxo(utxo.outpoint, psbt_input, foreign_utxo_satisfaction) .unwrap(); let psbt = builder.finish().unwrap(); let tx = psbt.extract_tx().expect("failed to extract tx"); @@ -1488,11 +1480,8 @@ fn test_add_foreign_utxo_invalid_psbt_input() { .unwrap(); let mut builder = wallet.build_tx(); - let result = builder.add_foreign_utxo( - outpoint, - psbt::Input::default(), - foreign_utxo_satisfaction.to_wu() as usize, - ); + let result = + builder.add_foreign_utxo(outpoint, psbt::Input::default(), foreign_utxo_satisfaction); assert!(matches!(result, Err(AddForeignUtxoError::MissingUtxo))); } @@ -1520,7 +1509,7 @@ fn test_add_foreign_utxo_where_outpoint_doesnt_match_psbt_input() { non_witness_utxo: Some(tx1.as_ref().clone()), ..Default::default() }, - satisfaction_weight.to_wu() as usize + satisfaction_weight ) .is_err(), "should fail when outpoint doesn't match psbt_input" @@ -1533,7 +1522,7 @@ fn test_add_foreign_utxo_where_outpoint_doesnt_match_psbt_input() { non_witness_utxo: Some(tx2.as_ref().clone()), ..Default::default() }, - satisfaction_weight.to_wu() as usize + satisfaction_weight ) .is_ok(), "should be ok when outpoint does match psbt_input" @@ -1565,11 +1554,7 @@ fn test_add_foreign_utxo_only_witness_utxo() { ..Default::default() }; builder - .add_foreign_utxo( - utxo2.outpoint, - psbt_input, - satisfaction_weight.to_wu() as usize, - ) + .add_foreign_utxo(utxo2.outpoint, psbt_input, satisfaction_weight) .unwrap(); assert!( builder.finish().is_err(), @@ -1585,11 +1570,7 @@ fn test_add_foreign_utxo_only_witness_utxo() { }; builder .only_witness_utxo() - .add_foreign_utxo( - utxo2.outpoint, - psbt_input, - satisfaction_weight.to_wu() as usize, - ) + .add_foreign_utxo(utxo2.outpoint, psbt_input, satisfaction_weight) .unwrap(); assert!( builder.finish().is_ok(), @@ -1605,11 +1586,7 @@ fn test_add_foreign_utxo_only_witness_utxo() { ..Default::default() }; builder - .add_foreign_utxo( - utxo2.outpoint, - psbt_input, - satisfaction_weight.to_wu() as usize, - ) + .add_foreign_utxo(utxo2.outpoint, psbt_input, satisfaction_weight) .unwrap(); assert!( builder.finish().is_ok(), @@ -3420,11 +3397,7 @@ fn test_taproot_foreign_utxo() { let mut builder = wallet1.build_tx(); builder .add_recipient(addr.script_pubkey(), Amount::from_sat(60_000)) - .add_foreign_utxo( - utxo.outpoint, - psbt_input, - foreign_utxo_satisfaction.to_wu() as usize, - ) + .add_foreign_utxo(utxo.outpoint, psbt_input, foreign_utxo_satisfaction) .unwrap(); let psbt = builder.finish().unwrap(); let sent_received =