From a0c58a4a8b4244d7c541906c61d1343dd6acdccd Mon Sep 17 00:00:00 2001
From: yancy <yancy@yancy.lol>
Date: Mon, 30 Sep 2024 18:55:28 -0500
Subject: [PATCH] Add checked weight division to Amount

---
 units/src/amount.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/units/src/amount.rs b/units/src/amount.rs
index ca2324845..4a59661b2 100644
--- a/units/src/amount.rs
+++ b/units/src/amount.rs
@@ -5,6 +5,9 @@
 //! This module mainly introduces the [`Amount`] and [`SignedAmount`] types.
 //! We refer to the documentation on the types for more information.
 
+#[cfg(feature = "alloc")]
+use crate::{Weight, FeeRate};
+
 #[cfg(feature = "alloc")]
 use alloc::string::{String, ToString};
 use core::cmp::Ordering;
@@ -1039,6 +1042,22 @@ impl Amount {
     /// Returns [`None`] if overflow occurred.
     pub fn checked_div(self, rhs: u64) -> Option<Amount> { self.0.checked_div(rhs).map(Amount) }
 
+    /// Checked weight division.
+    ///
+    /// Be aware that integer division loses the remainder if no exact division
+    /// can be made.  This method rounds up ensuring the transaction fee-rate is
+    /// sufficient.  If you wish to round-down, use the unchecked version instead.
+    ///
+    /// [`None`] is returned if an overflow occurred.
+    #[cfg(feature = "alloc")]
+    pub fn checked_div_by_weight(self, rhs: Weight) -> Option<FeeRate> {
+        let sats = self.0.checked_mul(1000)?;
+        let wu = rhs.to_wu();
+
+        let fee_rate = sats.checked_add(wu.checked_sub(1)?)?.checked_div(wu)?;
+        Some(FeeRate::from_sat_per_kwu(fee_rate))
+    }
+
     /// Checked remainder.
     ///
     /// Returns [`None`] if overflow occurred.
@@ -2219,6 +2238,31 @@ mod tests {
         assert_eq!(ssat(-6).checked_div(2), Some(ssat(-3)));
     }
 
+    #[cfg(feature = "alloc")]
+    #[test]
+    fn amount_checked_div_by_weight() {
+        let weight = Weight::from_kwu(1).unwrap();
+        let fee_rate = Amount::from_sat(1)
+            .checked_div_by_weight(weight)
+            .unwrap();
+        // 1 sats / 1,000 wu = 1 sats/kwu
+        assert_eq!(fee_rate, FeeRate::from_sat_per_kwu(1));
+
+        let weight = Weight::from_wu(381);
+        let fee_rate = Amount::from_sat(329)
+            .checked_div_by_weight(weight)
+            .unwrap();
+        // 329 sats / 381 wu = 863.5 sats/kwu
+        // round up to 864
+        assert_eq!(fee_rate, FeeRate::from_sat_per_kwu(864));
+
+        let fee_rate = Amount::MAX.checked_div_by_weight(weight);
+        assert!(fee_rate.is_none());
+
+        let fee_rate = Amount::ONE_SAT.checked_div_by_weight(Weight::ZERO);
+        assert!(fee_rate.is_none());
+    }
+
     #[test]
     #[cfg(not(debug_assertions))]
     fn unchecked_amount_add() {