Add fee_amount() and fee_rate() functions to PsbtUtils trait

Author: notmandatory

src/psbt/mod.rs

@@ -9,11 +9,20 @@
 // You may not use this file except in accordance with one or both of these
 // licenses.
 
+use crate::FeeRate;
 use bitcoin::util::psbt::PartiallySignedTransaction as Psbt;
 use bitcoin::TxOut;
 
 pub trait PsbtUtils {
     fn get_utxo_for(&self, input_index: usize) -> Option<TxOut>;
+
+    /// The total transaction fee amount, sum of input amounts minus sum of output amounts, in Sats.
+    fn fee_amount(&self) -> u64;
+
+    /// The transaction's fee rate. This value will only be accurate if calculated AFTER the
+    /// `PartiallySignedTransaction` is finalized and all witness/signature data is added to the
+    /// transaction.
+    fn fee_rate(&self) -> FeeRate;
 }
 
 impl PsbtUtils for Psbt {
@@ -37,15 +46,35 @@ impl PsbtUtils for Psbt {
             None
         }
     }
+
+    fn fee_amount(&self) -> u64 {
+        let input_amount = &self
+            .inputs
+            .iter()
+            .fold(0, |acc, i| acc + i.witness_utxo.as_ref().unwrap().value);
+        let output_amount = &self
+            .unsigned_tx
+            .output
+            .iter()
+            .fold(0, |acc, o| acc + o.value);
+        input_amount - output_amount
+    }
+
+    fn fee_rate(&self) -> FeeRate {
+        let fee_amount = self.fee_amount();
+        let weight = self.clone().extract_tx().weight();
+        FeeRate::from_wu(fee_amount, weight)
+    }
 }
 
 #[cfg(test)]
 mod test {
     use crate::bitcoin::TxIn;
     use crate::psbt::Psbt;
     use crate::wallet::AddressIndex;
+    use crate::wallet::AddressIndex::New;
     use crate::wallet::{get_funded_wallet, test::get_test_wpkh};
-    use crate::SignOptions;
+    use crate::{psbt, FeeRate, SignOptions};
     use std::str::FromStr;
 
     // from bip 174
@@ -118,4 +147,33 @@ mod test {
 
         let _ = wallet.sign(&mut psbt, SignOptions::default()).unwrap();
     }
+
+    #[test]
+    fn test_finalized_single_xprv_fee_rate() {
+        use psbt::PsbtUtils;
+
+        let expected_fee_rate = 1.2345;
+
+        let (wallet, _, _) = get_funded_wallet("wpkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)");
+        let addr = wallet.get_address(New).unwrap();
+        let mut builder = wallet.build_tx();
+        builder.drain_to(addr.script_pubkey()).drain_wallet();
+        builder.fee_rate(FeeRate::from_sat_per_vb(expected_fee_rate));
+        let (mut psbt, _) = builder.finish().unwrap();
+
+        let fee_amount = psbt.fee_amount();
+        dbg!(fee_amount);
+
+        let unfinalized_fee_rate = psbt.fee_rate();
+        dbg!(unfinalized_fee_rate.as_sat_vb());
+
+        let finalized = wallet.sign(&mut psbt, Default::default()).unwrap();
+        assert!(finalized);
+
+        let finalized_fee_rate = psbt.fee_rate();
+        dbg!(finalized_fee_rate.as_sat_vb());
+
+        assert!(finalized_fee_rate.as_sat_vb() >= expected_fee_rate);
+        assert!(finalized_fee_rate.as_sat_vb() < unfinalized_fee_rate.as_sat_vb());
+    }
 }