From 2d56909bd7921a5ad88ccd9898109883620b2ac2 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 | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/src/util/taproot.rs b/src/util/taproot.rs
index aaa2b0e99c..478803ea15 100644
--- a/src/util/taproot.rs
+++ b/src/util/taproot.rs
@@ -451,19 +451,18 @@ 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 current builder if the builder is not finalized. See [`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)) => {
+    ) -> Result<TaprootSpendInfo, TaprootBuilder> {
+        match (self.branch.len(), self.branch.pop()) {
+            (0, None) => Ok(TaprootSpendInfo::new_key_spend(secp, internal_key, None)),
+            (1, Some(Some(node))) => {
                 Ok(TaprootSpendInfo::from_node_info(secp, internal_key, node))
             }
-            _ => Err(TaprootBuilderError::IncompleteTree),
-
+            _ => Err(self),
         }
     }
 
@@ -1013,8 +1012,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 +1033,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 +1051,6 @@ impl std::error::Error for TaprootBuilderError {
             InvalidMerkleTreeDepth(_)
                 | NodeNotInDfsOrder
                 | OverCompleteTree
-                | IncompleteTree
                 | EmptyTree => None
         }
     }