refactor(wallet)!: change WeightedUtxo to use Weight type

oleonardolima · Jun 26, 2024 · 438cd46 · 438cd46
1 parent 275e069
commit 438cd46
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 68 deletions.
diff --git 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
@@ -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<WeightedUtxo> {
         (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
@@ -31,14 +31,17 @@ 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,
     Transaction, TxOut, Txid, Witness,
 };
 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,15 +1990,14 @@ 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;
                 (utxo, {
                     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
@@ -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
@@ -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 =