Add TaprootSpendData data structure, equivalent to script map for P2[W]SH

This data structures stores all information necessary for spending a taproot
output (the internal key, the Merkle root, and the control blocks for every
script leaf).

It is added to signing providers, and populated by the tr() descriptor.

5 files changed, 105 insertions, 7 deletions

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index 51cf8a7d62..84a8b06c5c 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -843,7 +843,9 @@ protected:
         XOnlyPubKey xpk(keys[0]);
         if (!xpk.IsFullyValid()) return {};
         builder.Finalize(xpk);
-        return Vector(GetScriptForDestination(builder.GetOutput()));
+        WitnessV1Taproot output = builder.GetOutput();
+        out.tr_spenddata[output].Merge(builder.GetSpendData());
+        return Vector(GetScriptForDestination(output));
     }
     bool ToStringSubScriptHelper(const SigningProvider* arg, std::string& ret, bool priv, bool normalized) const override
     {
diff --git a/src/script/signingprovider.cpp b/src/script/signingprovider.cpp
index 9781ec32af..b80fbe22ce 100644
--- a/src/script/signingprovider.cpp
+++ b/src/script/signingprovider.cpp
@@ -44,6 +44,11 @@ bool HidingSigningProvider::GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& inf
     return m_provider->GetKeyOrigin(keyid, info);
 }
 
+bool HidingSigningProvider::GetTaprootSpendData(const XOnlyPubKey& output_key, TaprootSpendData& spenddata) const
+{
+    return m_provider->GetTaprootSpendData(output_key, spenddata);
+}
+
 bool FlatSigningProvider::GetCScript(const CScriptID& scriptid, CScript& script) const { return LookupHelper(scripts, scriptid, script); }
 bool FlatSigningProvider::GetPubKey(const CKeyID& keyid, CPubKey& pubkey) const { return LookupHelper(pubkeys, keyid, pubkey); }
 bool FlatSigningProvider::GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& info) const
@@ -54,6 +59,10 @@ bool FlatSigningProvider::GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& info)
     return ret;
 }
 bool FlatSigningProvider::GetKey(const CKeyID& keyid, CKey& key) const { return LookupHelper(keys, keyid, key); }
+bool FlatSigningProvider::GetTaprootSpendData(const XOnlyPubKey& output_key, TaprootSpendData& spenddata) const
+{
+    return LookupHelper(tr_spenddata, output_key, spenddata);
+}
 
 FlatSigningProvider Merge(const FlatSigningProvider& a, const FlatSigningProvider& b)
 {
@@ -66,6 +75,10 @@ FlatSigningProvider Merge(const FlatSigningProvider& a, const FlatSigningProvide
     ret.keys.insert(b.keys.begin(), b.keys.end());
     ret.origins = a.origins;
     ret.origins.insert(b.origins.begin(), b.origins.end());
+    ret.tr_spenddata = a.tr_spenddata;
+    for (const auto& [output_key, spenddata] : b.tr_spenddata) {
+        ret.tr_spenddata[output_key].Merge(spenddata);
+    }
     return ret;
 }
 
diff --git a/src/script/signingprovider.h b/src/script/signingprovider.h
index 76f31d2f6f..939ae10622 100644
--- a/src/script/signingprovider.h
+++ b/src/script/signingprovider.h
@@ -25,6 +25,7 @@ public:
     virtual bool GetKey(const CKeyID &address, CKey& key) const { return false; }
     virtual bool HaveKey(const CKeyID &address) const { return false; }
     virtual bool GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& info) const { return false; }
+    virtual bool GetTaprootSpendData(const XOnlyPubKey& output_key, TaprootSpendData& spenddata) const { return false; }
 };
 
 extern const SigningProvider& DUMMY_SIGNING_PROVIDER;
@@ -42,6 +43,7 @@ public:
     bool GetPubKey(const CKeyID& keyid, CPubKey& pubkey) const override;
     bool GetKey(const CKeyID& keyid, CKey& key) const override;
     bool GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& info) const override;
+    bool GetTaprootSpendData(const XOnlyPubKey& output_key, TaprootSpendData& spenddata) const override;
 };
 
 struct FlatSigningProvider final : public SigningProvider
@@ -50,11 +52,13 @@ struct FlatSigningProvider final : public SigningProvider
     std::map<CKeyID, CPubKey> pubkeys;
     std::map<CKeyID, std::pair<CPubKey, KeyOriginInfo>> origins;
     std::map<CKeyID, CKey> keys;
+    std::map<XOnlyPubKey, TaprootSpendData> tr_spenddata; /** Map from output key to spend data. */
 
     bool GetCScript(const CScriptID& scriptid, CScript& script) const override;
     bool GetPubKey(const CKeyID& keyid, CPubKey& pubkey) const override;
     bool GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& info) const override;
     bool GetKey(const CKeyID& keyid, CKey& key) const override;
+    bool GetTaprootSpendData(const XOnlyPubKey& output_key, TaprootSpendData& spenddata) const override;
 };
 
 FlatSigningProvider Merge(const FlatSigningProvider& a, const FlatSigningProvider& b);
diff --git a/src/script/standard.cpp b/src/script/standard.cpp
index a4b11cc0a9..748f00dda5 100644
--- a/src/script/standard.cpp
+++ b/src/script/standard.cpp
@@ -377,6 +377,16 @@ bool IsValidDestination(const CTxDestination& dest) {
 /*static*/ TaprootBuilder::NodeInfo TaprootBuilder::Combine(NodeInfo&& a, NodeInfo&& b)
 {
     NodeInfo ret;
+    /* Iterate over all tracked leaves in a, add b's hash to their Merkle branch, and move them to ret. */
+    for (auto& leaf : a.leaves) {
+        leaf.merkle_branch.push_back(b.hash);
+        ret.leaves.emplace_back(std::move(leaf));
+    }
+    /* Iterate over all tracked leaves in b, add a's hash to their Merkle branch, and move them to ret. */
+    for (auto& leaf : b.leaves) {
+        leaf.merkle_branch.push_back(a.hash);
+        ret.leaves.emplace_back(std::move(leaf));
+    }
     /* 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();
@@ -386,6 +396,21 @@ bool IsValidDestination(const CTxDestination& dest) {
     return ret;
 }
 
+void TaprootSpendData::Merge(TaprootSpendData other)
+{
+    // TODO: figure out how to better deal with conflicting information
+    // being merged.
+    if (internal_key.IsNull() && !other.internal_key.IsNull()) {
+        internal_key = other.internal_key;
+    }
+    if (merkle_root.IsNull() && !other.merkle_root.IsNull()) {
+        merkle_root = other.merkle_root;
+    }
+    for (auto& [key, control_blocks] : other.scripts) {
+        scripts[key].merge(std::move(control_blocks));
+    }
+}
+
 void TaprootBuilder::Insert(TaprootBuilder::NodeInfo&& node, int depth)
 {
     assert(depth >= 0 && (size_t)depth <= TAPROOT_CONTROL_MAX_NODE_COUNT);
@@ -435,13 +460,14 @@ void TaprootBuilder::Insert(TaprootBuilder::NodeInfo&& node, int depth)
     return branch.size() == 0 || (branch.size() == 1 && branch[0]);
 }
 
-TaprootBuilder& TaprootBuilder::Add(int depth, const CScript& script, int leaf_version)
+TaprootBuilder& TaprootBuilder::Add(int depth, const CScript& script, int leaf_version, bool track)
 {
     assert((leaf_version & ~TAPROOT_LEAF_MASK) == 0);
     if (!IsValid()) return *this;
-    /* Construct NodeInfo object with leaf hash. */
+    /* Construct NodeInfo object with leaf hash and (if track is true) also leaf information. */
     NodeInfo node;
     node.hash = (CHashWriter{HASHER_TAPLEAF} << uint8_t(leaf_version) << script).GetSHA256();
+    if (track) node.leaves.emplace_back(LeafInfo{script, leaf_version, {}});
     /* Insert into the branch. */
     Insert(std::move(node), depth);
     return *this;
@@ -464,8 +490,33 @@ TaprootBuilder& TaprootBuilder::Finalize(const XOnlyPubKey& internal_key)
     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;
+    std::tie(m_output_key, m_parity) = *ret;
     return *this;
 }
 
 WitnessV1Taproot TaprootBuilder::GetOutput() { return WitnessV1Taproot{m_output_key}; }
+
+TaprootSpendData TaprootBuilder::GetSpendData() const
+{
+    TaprootSpendData spd;
+    spd.merkle_root = m_branch.size() == 0 ? uint256() : m_branch[0]->hash;
+    spd.internal_key = m_internal_key;
+    if (m_branch.size()) {
+        // If any script paths exist, they have been combined into the root m_branch[0]
+        // by now. Compute the control block for each of its tracked leaves, and put them in
+        // spd.scripts.
+        for (const auto& leaf : m_branch[0]->leaves) {
+            std::vector<unsigned char> control_block;
+            control_block.resize(TAPROOT_CONTROL_BASE_SIZE + TAPROOT_CONTROL_NODE_SIZE * leaf.merkle_branch.size());
+            control_block[0] = leaf.leaf_version | (m_parity ? 1 : 0);
+            std::copy(m_internal_key.begin(), m_internal_key.end(), control_block.begin() + 1);
+            if (leaf.merkle_branch.size()) {
+                std::copy(leaf.merkle_branch[0].begin(),
+                          leaf.merkle_branch[0].begin() + TAPROOT_CONTROL_NODE_SIZE * leaf.merkle_branch.size(),
+                          control_block.begin() + TAPROOT_CONTROL_BASE_SIZE);
+            }
+            spd.scripts[{leaf.script, leaf.leaf_version}].insert(std::move(control_block));
+        }
+    }
+    return spd;
+}
diff --git a/src/script/standard.h b/src/script/standard.h
index d7ea5cef27..8db17b2779 100644
--- a/src/script/standard.h
+++ b/src/script/standard.h
@@ -11,6 +11,7 @@
 #include <uint256.h>
 #include <util/hash_type.h>
 
+#include <map>
 #include <string>
 #include <variant>
 
@@ -209,15 +210,38 @@ CScript GetScriptForRawPubKey(const CPubKey& pubkey);
 /** Generate a multisig script. */
 CScript GetScriptForMultisig(int nRequired, const std::vector<CPubKey>& keys);
 
+struct TaprootSpendData
+{
+    /** The BIP341 internal key. */
+    XOnlyPubKey internal_key;
+    /** The Merkle root of the script tree (0 if no scripts). */
+    uint256 merkle_root;
+    /** Map from (script, leaf_version) to (sets of) control blocks. */
+    std::map<std::pair<CScript, int>, std::set<std::vector<unsigned char>>> scripts;
+    /** Merge other TaprootSpendData (for the same scriptPubKey) into this. */
+    void Merge(TaprootSpendData other);
+};
+
 /** Utility class to construct Taproot outputs from internal key and script tree. */
 class TaprootBuilder
 {
 private:
+    /** Information about a tracked leaf in the Merkle tree. */
+    struct LeafInfo
+    {
+        CScript script;                      //!< The script.
+        int leaf_version;                    //!< The leaf version for that script.
+        std::vector<uint256> merkle_branch;  //!< The hashing partners above this leaf.
+    };
+
     /** Information associated with a node in the Merkle tree. */
     struct NodeInfo
     {
         /** Merkle hash of this node. */
         uint256 hash;
+        /** Tracked leaves underneath this node (either from the node itself, or its children).
+         *  The merkle_branch field for each is the partners to get to *this* node. */
+        std::vector<LeafInfo> leaves;
     };
     /** Whether the builder is in a valid state so far. */
     bool m_valid = true;
@@ -260,7 +284,8 @@ private:
     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. */
+    XOnlyPubKey m_output_key;    //!< The output key, computed when finalizing.
+    bool m_parity;               //!< The tweak parity, computed when finalizing.
 
     /** Combine information about a parent Merkle tree node from its child nodes. */
     static NodeInfo Combine(NodeInfo&& a, NodeInfo&& b);
@@ -269,8 +294,9 @@ private:
 
 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);
+     *  in depth-first traversal order of binary tree. If track is true, it will be included in
+     *  the GetSpendData() output. */
+    TaprootBuilder& Add(int depth, const CScript& script, int leaf_version, bool track = true);
     /** 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.
@@ -285,6 +311,8 @@ public:
     WitnessV1Taproot GetOutput();
     /** Check if a list of depths is legal (will lead to IsComplete()). */
     static bool ValidDepths(const std::vector<int>& depths);
+    /** Compute spending data (after Finalize()). */
+    TaprootSpendData GetSpendData() const;
 };
 
 #endif // BITCOIN_SCRIPT_STANDARD_H