From 435eb9ba34cda0650bf2aac0192b9680c0f82a30 Mon Sep 17 00:00:00 2001
From: sanket1729 <sanket1729@gmail.com>
Date: Thu, 28 Jul 2022 17:20:16 -0700
Subject: [PATCH] Update TaprootBuilder::finalize

This commit does two things:
1) BUG FIX: We should have checked that the length of the branch is 1
before computing the spend_info on it. This was introduced in #936, but
slipped past my review :(
2) Update the return type to return the consumed `self` value. This
function can only error when there the tree building is not complete.
Returning TaprootBuilderError causes issues in downstream projects that
need to deal with error cases which cannot happen in this function
---
 src/util/taproot.rs | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/util/taproot.rs b/src/util/taproot.rs
index aaa2b0e99c..5c4cf0a1a8 100644
--- a/src/util/taproot.rs
+++ b/src/util/taproot.rs
@@ -199,7 +199,8 @@ impl TaprootSpendInfo {
         I: IntoIterator<Item=(u32, Script)>,
         C: secp256k1::Verification,
     {
-        TaprootBuilder::with_huffman_tree(script_weights)?.finalize(secp, internal_key)
+        let builder = TaprootBuilder::with_huffman_tree(script_weights)?;
+        Ok(builder.finalize(secp, internal_key).expect("Huffman Tree is always complete"))
     }
 
     /// Creates a new key spend with `internal_key` and `merkle_root`. Provide [`None`] for
@@ -451,19 +452,27 @@ impl TaprootBuilder {
 
     /// Creates a [`TaprootSpendInfo`] with the given internal key.
     ///
-    // TODO: in a future breaking API change, this no longer needs to return result
+    /// Returns the unmodified builder as Err if the builder is not finalized.
+    /// See also [`TaprootBuilder::is_finalized`]
     pub fn finalize<C: secp256k1::Verification>(
         mut self,
         secp: &Secp256k1<C>,
         internal_key: UntweakedPublicKey,
-    ) -> Result<TaprootSpendInfo, TaprootBuilderError> {
-        match self.branch.pop() {
-            None => Ok(TaprootSpendInfo::new_key_spend(secp, internal_key, None)),
-            Some(Some(node)) => {
-                Ok(TaprootSpendInfo::from_node_info(secp, internal_key, node))
+    ) -> Result<TaprootSpendInfo, TaprootBuilder> {
+        match self.branch.len() {
+            0 => Ok(TaprootSpendInfo::new_key_spend(secp, internal_key, None)),
+            1 => {
+                if let Some(Some(node)) = self.branch.pop() {
+                    Ok(TaprootSpendInfo::from_node_info(secp, internal_key, node))
+                } else {
+                    // This should be unreachable as
+                    // - The builder invariant guarantees that the last element is always Some
+                    // - The size of vector is checked to be 1
+                    // But still we return an error in order to avoid a panic reasoning.
+                    Err(self)
+                }
             }
-            _ => Err(TaprootBuilderError::IncompleteTree),
-
+            _ => Err(self),
         }
     }
 
@@ -1013,8 +1022,6 @@ pub enum TaprootBuilderError {
     OverCompleteTree,
     /// Invalid taproot internal key.
     InvalidInternalKey(secp256k1::Error),
-    /// Called finalize on an incomplete tree.
-    IncompleteTree,
     /// Called finalize on a empty tree.
     EmptyTree,
 }
@@ -1036,9 +1043,6 @@ impl fmt::Display for TaprootBuilderError {
             TaprootBuilderError::InvalidInternalKey(ref e) => {
                 write_err!(f, "invalid internal x-only key"; e)
             }
-            TaprootBuilderError::IncompleteTree => {
-                write!(f, "Called finalize on an incomplete tree")
-            }
             TaprootBuilderError::EmptyTree => {
                 write!(f, "Called finalize on an empty tree")
             }
@@ -1057,7 +1061,6 @@ impl std::error::Error for TaprootBuilderError {
             InvalidMerkleTreeDepth(_)
                 | NodeNotInDfsOrder
                 | OverCompleteTree
-                | IncompleteTree
                 | EmptyTree => None
         }
     }