Add TaprootBuilder class

This class functions as a utility for building taproot outputs, from
internal key and script leaves.

7 files changed, 262 insertions, 2 deletions

diff --git a/src/pubkey.cpp b/src/pubkey.cpp
index 8382897f95..51cc826b00 100644
--- a/src/pubkey.cpp
+++ b/src/pubkey.cpp
@@ -180,6 +180,12 @@ XOnlyPubKey::XOnlyPubKey(Span<const unsigned char> bytes)
     std::copy(bytes.begin(), bytes.end(), m_keydata.begin());
 }
 
+bool XOnlyPubKey::IsFullyValid() const
+{
+    secp256k1_xonly_pubkey pubkey;
+    return secp256k1_xonly_pubkey_parse(secp256k1_context_verify, &pubkey, m_keydata.data());
+}
+
 bool XOnlyPubKey::VerifySchnorr(const uint256& msg, Span<const unsigned char> sigbytes) const
 {
     assert(sigbytes.size() == 64);
diff --git a/src/pubkey.h b/src/pubkey.h
index b9d5f5d622..4de1807a7f 100644
--- a/src/pubkey.h
+++ b/src/pubkey.h
@@ -229,6 +229,11 @@ public:
     XOnlyPubKey(const XOnlyPubKey&) = default;
     XOnlyPubKey& operator=(const XOnlyPubKey&) = default;
 
+    /** Determine if this pubkey is fully valid. This is true for approximately 50% of all
+     *  possible 32-byte arrays. If false, VerifySchnorr and CreatePayToContract will always
+     *  fail. */
+    bool IsFullyValid() const;
+
     /** Construct an x-only pubkey from exactly 32 bytes. */
     explicit XOnlyPubKey(Span<const unsigned char> bytes);
 
diff --git a/src/script/interpreter.cpp b/src/script/interpreter.cpp
index 5f04d486b1..3c3c3ac1a8 100644
--- a/src/script/interpreter.cpp
+++ b/src/script/interpreter.cpp
@@ -1484,8 +1484,8 @@ template PrecomputedTransactionData::PrecomputedTransactionData(const CTransacti
 template PrecomputedTransactionData::PrecomputedTransactionData(const CMutableTransaction& txTo);
 
 static const CHashWriter HASHER_TAPSIGHASH = TaggedHash("TapSighash");
-static const CHashWriter HASHER_TAPLEAF = TaggedHash("TapLeaf");
-static const CHashWriter HASHER_TAPBRANCH = TaggedHash("TapBranch");
+const CHashWriter HASHER_TAPLEAF = TaggedHash("TapLeaf");
+const CHashWriter HASHER_TAPBRANCH = TaggedHash("TapBranch");
 
 static bool HandleMissingData(MissingDataBehavior mdb)
 {
diff --git a/src/script/interpreter.h b/src/script/interpreter.h
index 212de17c7b..fa4ee83e04 100644
--- a/src/script/interpreter.h
+++ b/src/script/interpreter.h
@@ -6,6 +6,7 @@
 #ifndef BITCOIN_SCRIPT_INTERPRETER_H
 #define BITCOIN_SCRIPT_INTERPRETER_H
 
+#include <hash.h>
 #include <script/script_error.h>
 #include <span.h>
 #include <primitives/transaction.h>
@@ -218,6 +219,9 @@ static constexpr size_t TAPROOT_CONTROL_NODE_SIZE = 32;
 static constexpr size_t TAPROOT_CONTROL_MAX_NODE_COUNT = 128;
 static constexpr size_t TAPROOT_CONTROL_MAX_SIZE = TAPROOT_CONTROL_BASE_SIZE + TAPROOT_CONTROL_NODE_SIZE * TAPROOT_CONTROL_MAX_NODE_COUNT;
 
+extern const CHashWriter HASHER_TAPLEAF;    //!< Hasher with tag "TapLeaf" pre-fed to it.
+extern const CHashWriter HASHER_TAPBRANCH;  //!< Hasher with tag "TapBranch" pre-fed to it.
+
 template <class T>
 uint256 SignatureHash(const CScript& scriptCode, const T& txTo, unsigned int nIn, int nHashType, const CAmount& amount, SigVersion sigversion, const PrecomputedTransactionData* cache = nullptr);
 
diff --git a/src/script/standard.cpp b/src/script/standard.cpp
index 0a233b550a..a4b11cc0a9 100644
--- a/src/script/standard.cpp
+++ b/src/script/standard.cpp
@@ -6,8 +6,11 @@
 #include <script/standard.h>
 
 #include <crypto/sha256.h>
+#include <hash.h>
 #include <pubkey.h>
+#include <script/interpreter.h>
 #include <script/script.h>
+#include <util/strencodings.h>
 
 #include <string>
 
@@ -370,3 +373,99 @@ CScript GetScriptForMultisig(int nRequired, const std::vector<CPubKey>& keys)
 bool IsValidDestination(const CTxDestination& dest) {
     return dest.index() != 0;
 }
+
+/*static*/ TaprootBuilder::NodeInfo TaprootBuilder::Combine(NodeInfo&& a, NodeInfo&& b)
+{
+    NodeInfo ret;
+    /* Lexicographically sort a and b's hash, and compute parent hash. */
+    if (a.hash < b.hash) {
+        ret.hash = (CHashWriter(HASHER_TAPBRANCH) << a.hash << b.hash).GetSHA256();
+    } else {
+        ret.hash = (CHashWriter(HASHER_TAPBRANCH) << b.hash << a.hash).GetSHA256();
+    }
+    return ret;
+}
+
+void TaprootBuilder::Insert(TaprootBuilder::NodeInfo&& node, int depth)
+{
+    assert(depth >= 0 && (size_t)depth <= TAPROOT_CONTROL_MAX_NODE_COUNT);
+    /* We cannot insert a leaf at a lower depth while a deeper branch is unfinished. Doing
+     * so would mean the Add() invocations do not correspond to a DFS traversal of a
+     * binary tree. */
+    if ((size_t)depth + 1 < m_branch.size()) {
+        m_valid = false;
+        return;
+    }
+    /* As long as an entry in the branch exists at the specified depth, combine it and propagate up.
+     * The 'node' variable is overwritten here with the newly combined node. */
+    while (m_valid && m_branch.size() > (size_t)depth && m_branch[depth].has_value()) {
+        node = Combine(std::move(node), std::move(*m_branch[depth]));
+        m_branch.pop_back();
+        if (depth == 0) m_valid = false; /* Can't propagate further up than the root */
+        --depth;
+    }
+    if (m_valid) {
+        /* Make sure the branch is big enough to place the new node. */
+        if (m_branch.size() <= (size_t)depth) m_branch.resize((size_t)depth + 1);
+        assert(!m_branch[depth].has_value());
+        m_branch[depth] = std::move(node);
+    }
+}
+
+/*static*/ bool TaprootBuilder::ValidDepths(const std::vector<int>& depths)
+{
+    std::vector<bool> branch;
+    for (int depth : depths) {
+        // This inner loop corresponds to effectively the same logic on branch
+        // as what Insert() performs on the m_branch variable. Instead of
+        // storing a NodeInfo object, just remember whether or not there is one
+        // at that depth.
+        if (depth < 0 || (size_t)depth > TAPROOT_CONTROL_MAX_NODE_COUNT) return false;
+        if ((size_t)depth + 1 < branch.size()) return false;
+        while (branch.size() > (size_t)depth && branch[depth]) {
+            branch.pop_back();
+            if (depth == 0) return false;
+            --depth;
+        }
+        if (branch.size() <= (size_t)depth) branch.resize((size_t)depth + 1);
+        assert(!branch[depth]);
+        branch[depth] = true;
+    }
+    // And this check corresponds to the IsComplete() check on m_branch.
+    return branch.size() == 0 || (branch.size() == 1 && branch[0]);
+}
+
+TaprootBuilder& TaprootBuilder::Add(int depth, const CScript& script, int leaf_version)
+{
+    assert((leaf_version & ~TAPROOT_LEAF_MASK) == 0);
+    if (!IsValid()) return *this;
+    /* Construct NodeInfo object with leaf hash. */
+    NodeInfo node;
+    node.hash = (CHashWriter{HASHER_TAPLEAF} << uint8_t(leaf_version) << script).GetSHA256();
+    /* Insert into the branch. */
+    Insert(std::move(node), depth);
+    return *this;
+}
+
+TaprootBuilder& TaprootBuilder::AddOmitted(int depth, const uint256& hash)
+{
+    if (!IsValid()) return *this;
+    /* Construct NodeInfo object with the hash directly, and insert it into the branch. */
+    NodeInfo node;
+    node.hash = hash;
+    Insert(std::move(node), depth);
+    return *this;
+}
+
+TaprootBuilder& TaprootBuilder::Finalize(const XOnlyPubKey& internal_key)
+{
+    /* Can only call this function when IsComplete() is true. */
+    assert(IsComplete());
+    m_internal_key = internal_key;
+    auto ret = m_internal_key.CreateTapTweak(m_branch.size() == 0 ? nullptr : &m_branch[0]->hash);
+    assert(ret.has_value());
+    std::tie(m_output_key, std::ignore) = *ret;
+    return *this;
+}
+
+WitnessV1Taproot TaprootBuilder::GetOutput() { return WitnessV1Taproot{m_output_key}; }
diff --git a/src/script/standard.h b/src/script/standard.h
index a96b096fa7..d7ea5cef27 100644
--- a/src/script/standard.h
+++ b/src/script/standard.h
@@ -209,4 +209,82 @@ CScript GetScriptForRawPubKey(const CPubKey& pubkey);
 /** Generate a multisig script. */
 CScript GetScriptForMultisig(int nRequired, const std::vector<CPubKey>& keys);
 
+/** Utility class to construct Taproot outputs from internal key and script tree. */
+class TaprootBuilder
+{
+private:
+    /** Information associated with a node in the Merkle tree. */
+    struct NodeInfo
+    {
+        /** Merkle hash of this node. */
+        uint256 hash;
+    };
+    /** Whether the builder is in a valid state so far. */
+    bool m_valid = true;
+
+    /** The current state of the builder.
+     *
+     * For each level in the tree, one NodeInfo object may be present. m_branch[0]
+     * is information about the root; further values are for deeper subtrees being
+     * explored.
+     *
+     * For every right branch taken to reach the position we're currently
+     * working in, there will be a (non-nullopt) entry in m_branch corresponding
+     * to the left branch at that level.
+     *
+     * For example, imagine this tree:     - N0 -
+     *                                    /      \
+     *                                   N1      N2
+     *                                  /  \    /  \
+     *                                 A    B  C   N3
+     *                                            /  \
+     *                                           D    E
+     *
+     * Initially, m_branch is empty. After processing leaf A, it would become
+     * {nullopt, nullopt, A}. When processing leaf B, an entry at level 2 already
+     * exists, and it would thus be combined with it to produce a level 1 one,
+     * resulting in {nullopt, N1}. Adding C and D takes us to {nullopt, N1, C}
+     * and {nullopt, N1, C, D} respectively. When E is processed, it is combined
+     * with D, and then C, and then N1, to produce the root, resulting in {N0}.
+     *
+     * This structure allows processing with just O(log n) overhead if the leaves
+     * are computed on the fly.
+     *
+     * As an invariant, there can never be nullopt entries at the end. There can
+     * also not be more than 128 entries (as that would mean more than 128 levels
+     * in the tree). The depth of newly added entries will always be at least
+     * equal to the current size of m_branch (otherwise it does not correspond
+     * to a depth-first traversal of a tree). m_branch is only empty if no entries
+     * have ever be processed. m_branch having length 1 corresponds to being done.
+     */
+    std::vector<std::optional<NodeInfo>> m_branch;
+
+    XOnlyPubKey m_internal_key;  //!< The internal key, set when finalizing.
+    XOnlyPubKey m_output_key; //!< The output key, computed when finalizing. */
+
+    /** Combine information about a parent Merkle tree node from its child nodes. */
+    static NodeInfo Combine(NodeInfo&& a, NodeInfo&& b);
+    /** Insert information about a node at a certain depth, and propagate information up. */
+    void Insert(NodeInfo&& node, int depth);
+
+public:
+    /** Add a new script at a certain depth in the tree. Add() operations must be called
+     *  in depth-first traversal order of binary tree. */
+    TaprootBuilder& Add(int depth, const CScript& script, int leaf_version);
+    /** Like Add(), but for a Merkle node with a given hash to the tree. */
+    TaprootBuilder& AddOmitted(int depth, const uint256& hash);
+    /** Finalize the construction. Can only be called when IsComplete() is true.
+        internal_key.IsFullyValid() must be true. */
+    TaprootBuilder& Finalize(const XOnlyPubKey& internal_key);
+
+    /** Return true if so far all input was valid. */
+    bool IsValid() const { return m_valid; }
+    /** Return whether there were either no leaves, or the leaves form a Huffman tree. */
+    bool IsComplete() const { return m_valid && (m_branch.size() == 0 || (m_branch.size() == 1 && m_branch[0].has_value())); }
+    /** Compute scriptPubKey (after Finalize()). */
+    WitnessV1Taproot GetOutput();
+    /** Check if a list of depths is legal (will lead to IsComplete()). */
+    static bool ValidDepths(const std::vector<int>& depths);
+};
+
 #endif // BITCOIN_SCRIPT_STANDARD_H
diff --git a/src/test/script_standard_tests.cpp b/src/test/script_standard_tests.cpp
index 950f3b970a..a01d3fa03a 100644
--- a/src/test/script_standard_tests.cpp
+++ b/src/test/script_standard_tests.cpp
@@ -3,10 +3,12 @@
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
 #include <key.h>
+#include <key_io.h>
 #include <script/script.h>
 #include <script/signingprovider.h>
 #include <script/standard.h>
 #include <test/util/setup_common.h>
+#include <util/strencodings.h>
 
 #include <boost/test/unit_test.hpp>
 
@@ -378,4 +380,70 @@ BOOST_AUTO_TEST_CASE(script_standard_GetScriptFor_)
     BOOST_CHECK(result == expected);
 }
 
+BOOST_AUTO_TEST_CASE(script_standard_taproot_builder)
+{
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,1}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,0,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,0,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,0,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,1,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,1,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,1,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,2,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,2,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({0,2,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,0,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,0,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,0,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,1,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,1,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,1,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,2,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,2,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({1,2,2}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,0,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,0,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,0,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,1,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,1,1}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,1,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,2,0}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,2,1}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,2,2}), false);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({2,2,2,3,4,5,6,7,8,9,10,11,12,14,14,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,31,31,31,31,31,31,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,128}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({128,128,127,126,125,124,123,122,121,120,119,118,117,116,115,114,113,112,111,110,109,108,107,106,105,104,103,102,101,100,99,98,97,96,95,94,93,92,91,90,89,88,87,86,85,84,83,82,81,80,79,78,77,76,75,74,73,72,71,70,69,68,67,66,65,64,63,62,61,60,59,58,57,56,55,54,53,52,51,50,49,48,47,46,45,44,43,42,41,40,39,38,37,36,35,34,33,32,31,30,29,28,27,26,25,24,23,22,21,20,19,18,17,16,15,14,13,12,11,10,9,8,7,6,5,4,3,2,1}), true);
+    BOOST_CHECK_EQUAL(TaprootBuilder::ValidDepths({129,129,128,127,126,125,124,123,122,121,120,119,118,117,116,115,114,113,112,111,110,109,108,107,106,105,104,103,102,101,100,99,98,97,96,95,94,93,92,91,90,89,88,87,86,85,84,83,82,81,80,79,78,77,76,75,74,73,72,71,70,69,68,67,66,65,64,63,62,61,60,59,58,57,56,55,54,53,52,51,50,49,48,47,46,45,44,43,42,41,40,39,38,37,36,35,34,33,32,31,30,29,28,27,26,25,24,23,22,21,20,19,18,17,16,15,14,13,12,11,10,9,8,7,6,5,4,3,2,1}), false);
+
+    XOnlyPubKey key_inner{ParseHex("79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798")};
+    XOnlyPubKey key_1{ParseHex("c6047f9441ed7d6d3045406e95c07cd85c778e4b8cef3ca7abac09b95c709ee5")};
+    XOnlyPubKey key_2{ParseHex("f9308a019258c31049344f85f89d5229b531c845836f99b08601f113bce036f9")};
+    CScript script_1 = CScript() << ToByteVector(key_1) << OP_CHECKSIG;
+    CScript script_2 = CScript() << ToByteVector(key_2) << OP_CHECKSIG;
+    uint256 hash_3 = uint256S("31fe7061656bea2a36aa60a2f7ef940578049273746935d296426dc0afd86b68");
+
+    TaprootBuilder builder;
+    BOOST_CHECK(builder.IsValid() && builder.IsComplete());
+    builder.Add(2, script_2, 0xc0);
+    BOOST_CHECK(builder.IsValid() && !builder.IsComplete());
+    builder.AddOmitted(2, hash_3);
+    BOOST_CHECK(builder.IsValid() && !builder.IsComplete());
+    builder.Add(1, script_1, 0xc0);
+    BOOST_CHECK(builder.IsValid() && builder.IsComplete());
+    builder.Finalize(key_inner);
+    BOOST_CHECK(builder.IsValid() && builder.IsComplete());
+    BOOST_CHECK_EQUAL(EncodeDestination(builder.GetOutput()), "bc1pj6gaw944fy0xpmzzu45ugqde4rz7mqj5kj0tg8kmr5f0pjq8vnaqgynnge");
+}
+
 BOOST_AUTO_TEST_SUITE_END()