Move with_huffman_tree logic to TaprootBuilder

src/util/taproot.rs

@@ -191,24 +191,9 @@ pub struct TaprootSpendInfo {
 }
 
 impl TaprootSpendInfo {
-    /// Create a new [`TaprootSpendInfo`] from a list of script(with default script version) and
+    /// Create a new [`TaprootSpendInfo`] from a list of script(with default script version).
-    /// weights of satisfaction for that script. The weights represent the probability of
-    /// each branch being taken. If probabilities/weights for each condition are known,
-    /// constructing the tree as a Huffman tree is the optimal way to minimize average
-    /// case satisfaction cost. This function takes input an iterator of tuple(u64, &Script)
-    /// where usize represents the satisfaction weights of the branch.
-    /// For example, [(3, S1), (2, S2), (5, S3)] would construct a TapTree that has optimal
-    /// satisfaction weight when probability for S1 is 30%, S2 is 20% and S3 is 50%.
     ///
-    /// # Errors:
+    /// See [`TaprootBuilder::with_huffman_tree`] for more detailed documentation
-    ///
-    /// - When the optimal huffman tree has a depth more than 128
-    /// - If the provided list of script weights is empty
-    ///
-    /// # Edge Cases:
-    /// - If the script weight calculations overflow, a sub-optimal tree may be generated. This
-    ///   should not happen unless you are dealing with billions of branches with weights close to
-    ///   2^32.
     pub fn with_huffman_tree<C, I>(
         secp: &Secp256k1<C>,
         internal_key: UntweakedPublicKey,
@@ -218,29 +203,7 @@ impl TaprootSpendInfo {
         I: IntoIterator<Item=(u32, Script)>,
         C: secp256k1::Verification,
     {
-        let mut node_weights = BinaryHeap::<(Reverse<u64>, NodeInfo)>::new();
+        TaprootBuilder::with_huffman_tree(script_weights)?.finalize(secp, internal_key)
-        for (p, leaf) in script_weights {
-            node_weights.push((Reverse(p as u64), NodeInfo::new_leaf_with_ver(leaf, LeafVersion::TapScript)));
-        }
-        if node_weights.is_empty() {
-            return Err(TaprootBuilderError::IncompleteTree);
-        }
-        while node_weights.len() > 1 {
-            // Combine the last two elements and insert a new node
-            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
-            // N.B.: p1 + p2 can not practically saturate as you would need to have 2**32 max u32s
-            // 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 = Reverse(p1.0.saturating_add(p2.0));
-            node_weights.push((p, NodeInfo::combine(s1, s2)?));
-        }
-        // Every iteration of the loop reduces the node_weights.len() by exactly 1
-        // Therefore, the loop will eventually terminate with exactly 1 element
-        debug_assert!(node_weights.len() == 1);
-        let node = node_weights.pop().expect("huffman tree algorithm is broken").1;
-        return Ok(Self::from_node_info(secp, internal_key, node));
     }
 
     /// Create a new key spend with internal key and proided merkle root.
@@ -396,6 +359,56 @@ impl TaprootBuilder {
     pub fn new() -> Self {
         TaprootBuilder { branch: vec![] }
     }
+
+    /// Create a new [`TaprootBuilder`] from a list of script(with default script version) and
+    /// weights of satisfaction for that script. The weights represent the probability of
+    /// each branch being taken. If probabilities/weights for each condition are known,
+    /// constructing the tree as a Huffman tree is the optimal way to minimize average
+    /// case satisfaction cost. This function takes input an iterator of tuple(u64, &Script)
+    /// where usize represents the satisfaction weights of the branch.
+    /// For example, [(3, S1), (2, S2), (5, S3)] would construct a TapTree that has optimal
+    /// satisfaction weight when probability for S1 is 30%, S2 is 20% and S3 is 50%.
+    ///
+    /// # Errors:
+    ///
+    /// - When the optimal huffman tree has a depth more than 128
+    /// - If the provided list of script weights is empty
+    ///
+    /// # Edge Cases:
+    /// - If the script weight calculations overflow, a sub-optimal tree may be generated. This
+    ///   should not happen unless you are dealing with billions of branches with weights close to
+    ///   2^32.
+    pub fn with_huffman_tree<I>(
+        script_weights: I,
+    ) -> Result<Self, TaprootBuilderError>
+    where
+        I: IntoIterator<Item=(u32, Script)>,
+    {
+        let mut node_weights = BinaryHeap::<(Reverse<u64>, NodeInfo)>::new();
+        for (p, leaf) in script_weights {
+            node_weights.push((Reverse(p as u64), NodeInfo::new_leaf_with_ver(leaf, LeafVersion::TapScript)));
+        }
+        if node_weights.is_empty() {
+            return Err(TaprootBuilderError::IncompleteTree);
+        }
+        while node_weights.len() > 1 {
+            // Combine the last two elements and insert a new node
+            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
+            // N.B.: p1 + p2 can not practically saturate as you would need to have 2**32 max u32s
+            // 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 = Reverse(p1.0.saturating_add(p2.0));
+            node_weights.push((p, NodeInfo::combine(s1, s2)?));
+        }
+        // Every iteration of the loop reduces the node_weights.len() by exactly 1
+        // Therefore, the loop will eventually terminate with exactly 1 element
+        debug_assert!(node_weights.len() == 1);
+        let node = node_weights.pop().expect("huffman tree algorithm is broken").1;
+        Ok(TaprootBuilder{branch: vec![Some(node)]})
+    }
+
     /// Just like [`TaprootBuilder::add_leaf`] but allows to specify script version
     pub fn add_leaf_with_ver(
         self,