From cccd75d00443db0d6b7a35caf597d199c9fd897f Mon Sep 17 00:00:00 2001
From: Jeremy Rubin <j@rubin.io>
Date: Sun, 14 Nov 2021 17:55:30 -0800
Subject: [PATCH] Fix Weighting Addition to never error on overflow + prevent
 overflows from ever happening with wider integers

---
 src/util/taproot.rs | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/src/util/taproot.rs b/src/util/taproot.rs
index 72754c75..e5adebd2 100644
--- a/src/util/taproot.rs
+++ b/src/util/taproot.rs
@@ -208,9 +208,9 @@ impl TaprootSpendInfo {
         I: IntoIterator<Item = (u64, Script)>,
         C: secp256k1::Verification,
     {
-        let mut node_weights = BinaryHeap::<(u64, NodeInfo)>::new();
+        let mut node_weights = BinaryHeap::<(u128, NodeInfo)>::new();
         for (p, leaf) in script_weights {
-            node_weights.push((p, NodeInfo::new_leaf_with_ver(leaf, LeafVersion::default())));
+            node_weights.push((p as u128, NodeInfo::new_leaf_with_ver(leaf, LeafVersion::default())));
         }
         if node_weights.is_empty() {
             return Err(TaprootBuilderError::IncompleteTree);
@@ -220,7 +220,10 @@ impl TaprootSpendInfo {
             let (p1, s1) = node_weights.pop().expect("len must be at least two");
             let (p2, s2) = node_weights.pop().expect("len must be at least two");
             // Insert the sum of first two in the tree as a new node
-            let p = p1.checked_add(p2).ok_or(TaprootBuilderError::ScriptWeightOverflow)?;
+            // N.B.: p1 + p2 can never actually saturate as you would need to have 2**64 max u64s
+            // from the input to overflow. However, saturating is a reasonable behavior here as
+            // huffman tree construction would treat all such elements as "very likely".
+            let p = p1.saturating_add(p2);
             node_weights.push((p, NodeInfo::combine(s1, s2)?));
         }
         // Every iteration of the loop reduces the node_weights.len() by exactly 1
@@ -801,8 +804,6 @@ pub enum TaprootBuilderError {
     IncompleteTree,
     /// Called finalize on a empty tree
     EmptyTree,
-    /// Script weight overflow
-    ScriptWeightOverflow,
 }
 
 impl fmt::Display for TaprootBuilderError {
@@ -830,9 +831,6 @@ impl fmt::Display for TaprootBuilderError {
             TaprootBuilderError::EmptyTree => {
                 write!(f, "Called finalize on an empty tree")
             }
-            TaprootBuilderError::ScriptWeightOverflow => {
-                write!(f, "Script weight overflow in Huffman tree construction")
-            },
         }
     }
 }