fixup! txscript: add AssembleTaprootScriptTree func for creating inpu…

…t witnesses

txscript/taproot.go

@@ -7,7 +7,6 @@ package txscript
 import (
 	"bytes"
 	"fmt"
-	"math"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
@@ -362,6 +361,14 @@ func VerifyTaprootLeafCommitment(controlBlock *ControlBlock,
 		return fmt.Errorf("invalid witness commitment")
 	}
 
+	// Finally, we'll verify that the parity of the y coordinate of the
+	// public key we've derived matches the control block.
+	derivedYIsOdd := (taprootKey.SerializeCompressed()[0] ==
+		secp.PubKeyFormatCompressedOdd)
+	if controlBlock.OutputKeyYIsOdd != derivedYIsOdd {
+		return fmt.Errorf("invalid witness commitment")
+	}
+
 	// Otherwise, if we reach here, the commitment opening is valid and
 	// execution can continue.
 	return nil
@@ -427,6 +434,9 @@ func NewTapLeaf(leafVersion TapscriptLeafVersion, script []byte) TapLeaf {
 // digest is computed as: h_tapleaf(leafVersion || compactSizeof(script) ||
 // script).
 func (t TapLeaf) TapHash() chainhash.Hash {
+	// TODO(roasbeef): cache these and the branch due to the recursive
+	// call, so memoize
+
 	// The leaf encoding is: leafVersion || compactSizeof(script) ||
 	// script, where compactSizeof returns the compact size needed to
 	// encode the value.
@@ -481,8 +491,7 @@ func (t TapBranch) TapHash() chainhash.Hash {
 // hashes them into a branch. See The TapBranch method for the specifics.
 func tapBranchHash(l, r []byte) chainhash.Hash {
 	if bytes.Compare(l[:], r[:]) > 0 {
-		l = r
-		r = l
+		l, r = r, l
 	}
 
 	return *chainhash.TaggedHash(
@@ -496,6 +505,10 @@ type TapscriptProof struct {
 	// TapLeaf is the leaf that we want to prove inclusion for.
 	TapLeaf
 
+	// RootNode is the root of the tapscript tree, this will be used to
+	// compute what the final output key looks like.
+	RootNode TapNode
+
 	// InclusionProof is the tail end of the control block that contains
 	// the series of hashes (the sibling hashes up the tree), that when
 	// hashed together allow us to re-derive the top level taproot output.
@@ -505,10 +518,27 @@ type TapscriptProof struct {
 // ToControlBlock maps the tapscript proof into a fully valid control block
 // that can be used as a witness item for a tapscript spend.
 func (t *TapscriptProof) ToControlBlock(internalKey *btcec.PublicKey) ControlBlock {
+	// Compute the total level output commitment based on the populated
+	// root node.
+	rootHash := t.RootNode.TapHash()
+	taprootKey := ComputeTaprootOutputKey(
+		internalKey, rootHash[:],
+	)
+
+	// With the commitment computed we can obtain the bit that denotes if
+	// the resulting key has an odd y coordinate or not.
+	var outputKeyYIsOdd bool
+	if taprootKey.SerializeCompressed()[0] ==
+		secp.PubKeyFormatCompressedOdd {
+
+		outputKeyYIsOdd = true
+	}
+
 	return ControlBlock{
-		InternalKey:    internalKey,
-		LeafVersion:    t.TapLeaf.LeafVersion,
-		InclusionProof: t.InclusionProof,
+		InternalKey:     internalKey,
+		OutputKeyYIsOdd: outputKeyYIsOdd,
+		LeafVersion:     t.TapLeaf.LeafVersion,
+		InclusionProof:  t.InclusionProof,
 	}
 }
 
@@ -525,7 +555,7 @@ type IndexedTapScriptTree struct {
 	// LeafMerkleProofs is a slice that houses the series of merkle
 	// inclusion proofs for each leaf based on the input order of the
 	// leaves.
-	LeafMerkleProofs []*TapscriptProof
+	LeafMerkleProofs []TapscriptProof
 
 	// LeafProofIndex maps the TapHash() of a given leaf node to the index
 	// within the LeafMerkleProofs array above. This can be used to
@@ -538,25 +568,11 @@ type IndexedTapScriptTree struct {
 // space to hold information for the specified amount of leaves.
 func NewIndexedTapScriptTree(numLeaves int) *IndexedTapScriptTree {
 	return &IndexedTapScriptTree{
-		LeafMerkleProofs: make([]*TapscriptProof, numLeaves),
+		LeafMerkleProofs: make([]TapscriptProof, numLeaves),
 		LeafProofIndex:   make(map[chainhash.Hash]int, numLeaves),
 	}
 }
 
-// nextPowerOfTwo returns the next highest power of two from a given number if
-// it is not already a power of two.  This is a helper function used during the
-// calculation of a merkle tree.
-func nextPowerOfTwo(n int) int {
-	// Return the number if it's already a power of 2.
-	if n&(n-1) == 0 {
-		return n
-	}
-
-	// Figure out and return the next power of two.
-	exponent := uint(math.Log2(float64(n))) + 1
-	return 1 << exponent // 2^exponent
-}
-
 // hashTapNodes takes a left and right now, and returns the left and right tap
 // hashes, along with the new combined node. If both nodes are nil, nil
 // pointers are returned. If the right now is nil, then the left node is passed
@@ -606,40 +622,74 @@ func leafDescendants(node TapNode) []TapNode {
 // and an array-based representation are used to both generate the tree and
 // also accumulate all the necessary inclusion proofs in the same path. See the
 // comment of blockchain.BuildMerkleTreeStore for further details.
-func AssembleTaprootScriptTree(leaves ...TapLeaf) (*IndexedTapScriptTree, error) {
+func AssembleTaprootScriptTree(leaves ...TapLeaf) *IndexedTapScriptTree {
 	// If there's only a single leaf, then that becomes our root.
 	if len(leaves) == 1 {
 		// A lone leaf has no additional inclusion proof, as a verifier
 		// will just hash the leaf as the sole branch.
+		leaf := leaves[0]
 		return &IndexedTapScriptTree{
-			RootNode: leaves[0],
-			LeafMerkleProofs: []*TapscriptProof{
-				&TapscriptProof{
-					TapLeaf:        leaves[0],
+			RootNode: leaf,
+			LeafProofIndex: map[chainhash.Hash]int{
+				leaf.TapHash(): 0,
+			},
+			LeafMerkleProofs: []TapscriptProof{
+				{
+					TapLeaf:        leaf,
+					RootNode:       leaf,
 					InclusionProof: nil,
 				},
 			},
-		}, nil
+		}
 	}
 
-	var branches []TapBranch
+	// We'll start out by populating the leaf index which maps a leave's
+	// taphash to its index within the tree.
 	scriptTree := NewIndexedTapScriptTree(len(leaves))
-	for i := 0; i < len(leaves)-1; i += 2 {
+	for i, leaf := range leaves {
+		leafHash := leaf.TapHash()
+		scriptTree.LeafProofIndex[leafHash] = i
+	}
+
+	var branches []TapBranch
+	for i := 0; i < len(leaves); i += 2 {
 		// If there's only a single leaf left, then we'll merge this
 		// with the last branch we have.
 		if i == len(leaves)-1 {
 			branchToMerge := branches[len(branches)-1]
-			newBranch := NewTapBranch(branchToMerge, leaves[i])
+			leaf := leaves[i]
+			newBranch := NewTapBranch(branchToMerge, leaf)
 
 			branches[len(branches)-1] = newBranch
 
 			// The leaf includes the existing branch within its
-			// inclusion prof.
+			// inclusion proof.
 			branchHash := branchToMerge.TapHash()
+
+			scriptTree.LeafMerkleProofs[i].TapLeaf = leaf
 			scriptTree.LeafMerkleProofs[i].InclusionProof = append(
 				scriptTree.LeafMerkleProofs[i].InclusionProof,
 				branchHash[:]...,
 			)
+
+			// We'll also add this right hash to the inclusion of
+			// the left and right nodes of the branch.
+			lastLeafHash := leaf.TapHash()
+
+			leftLeafHash := branchToMerge.Left().TapHash()
+			leftLeafIndex := scriptTree.LeafProofIndex[leftLeafHash]
+			scriptTree.LeafMerkleProofs[leftLeafIndex].InclusionProof = append(
+				scriptTree.LeafMerkleProofs[leftLeafIndex].InclusionProof,
+				lastLeafHash[:]...,
+			)
+
+			rightLeafHash := branchToMerge.Right().TapHash()
+			rightLeafIndex := scriptTree.LeafProofIndex[rightLeafHash]
+			scriptTree.LeafMerkleProofs[rightLeafIndex].InclusionProof = append(
+				scriptTree.LeafMerkleProofs[rightLeafIndex].InclusionProof,
+				lastLeafHash[:]...,
+			)
+
 			continue
 		}
 
@@ -654,10 +704,13 @@ func AssembleTaprootScriptTree(leaves ...TapLeaf) (*IndexedTapScriptTree, error)
 		leftHash := left.TapHash()
 		rightHash := right.TapHash()
 
+		scriptTree.LeafMerkleProofs[i].TapLeaf = left
 		scriptTree.LeafMerkleProofs[i].InclusionProof = append(
 			scriptTree.LeafMerkleProofs[i].InclusionProof,
 			rightHash[:]...,
 		)
+
+		scriptTree.LeafMerkleProofs[i+1].TapLeaf = right
 		scriptTree.LeafMerkleProofs[i+1].InclusionProof = append(
 			scriptTree.LeafMerkleProofs[i+1].InclusionProof,
 			leftHash[:]...,
@@ -666,11 +719,12 @@ func AssembleTaprootScriptTree(leaves ...TapLeaf) (*IndexedTapScriptTree, error)
 
 	// In this second phase, we'll merge all the leaf branches we have one
 	// by one until we have our final root.
+	var rootNode TapNode
 	for len(branches) != 0 {
 		// When we only have a single branch left, then that becomes
 		// our root.
 		if len(branches) == 1 {
-			scriptTree.RootNode = branches[0]
+			rootNode = branches[0]
 			break
 		}
 
@@ -687,29 +741,38 @@ func AssembleTaprootScriptTree(leaves ...TapLeaf) (*IndexedTapScriptTree, error)
 		leftLeafDescendants := leafDescendants(left)
 		rightLeafDescendants := leafDescendants(right)
 
-		// For each left hash that's a leaf decedents, well add the
+		leftHash, rightHash := left.TapHash(), right.TapHash()
+
+		// For each left hash that's a leaf descendants, well add the
 		// right sibling as that sibling is needed to construct the new
 		// internal branch we just created. We also do the same for the
 		// siblings of the right node.
 		for _, leftLeaf := range leftLeafDescendants {
 			leafHash := leftLeaf.TapHash()
-
 			leafIndex := scriptTree.LeafProofIndex[leafHash]
+
 			scriptTree.LeafMerkleProofs[leafIndex].InclusionProof = append(
 				scriptTree.LeafMerkleProofs[leafIndex].InclusionProof,
-				leafHash[:]...,
+				rightHash[:]...,
 			)
 		}
 		for _, rightLeaf := range rightLeafDescendants {
 			leafHash := rightLeaf.TapHash()
-
 			leafIndex := scriptTree.LeafProofIndex[leafHash]
+
 			scriptTree.LeafMerkleProofs[leafIndex].InclusionProof = append(
 				scriptTree.LeafMerkleProofs[leafIndex].InclusionProof,
-				leafHash[:]...,
+				leftHash[:]...,
 			)
 		}
 	}
 
-	return scriptTree, nil
+	// Populate the top level root node pointer, as well as the pointer in
+	// each proof.
+	scriptTree.RootNode = rootNode
+	for i := range scriptTree.LeafMerkleProofs {
+		scriptTree.LeafMerkleProofs[i].RootNode = rootNode
+	}
+
+	return scriptTree
 }