Taproot descriptor inference

3 files changed, 210 insertions, 9 deletions

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index 84a8b06c5c..be97a618f3 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -244,7 +244,7 @@ class ConstPubkeyProvider final : public PubkeyProvider
     bool m_xonly;
 
 public:
-    ConstPubkeyProvider(uint32_t exp_index, const CPubKey& pubkey, bool xonly = false) : PubkeyProvider(exp_index), m_pubkey(pubkey), m_xonly(xonly) {}
+    ConstPubkeyProvider(uint32_t exp_index, const CPubKey& pubkey, bool xonly) : PubkeyProvider(exp_index), m_pubkey(pubkey), m_xonly(xonly) {}
     bool GetPubKey(int pos, const SigningProvider& arg, CPubKey& key, KeyOriginInfo& info, const DescriptorCache* read_cache = nullptr, DescriptorCache* write_cache = nullptr) override
     {
         key = m_pubkey;
@@ -931,7 +931,7 @@ std::unique_ptr<PubkeyProvider> ParsePubkeyInner(uint32_t key_exp_index, const S
             CPubKey pubkey(data);
             if (pubkey.IsFullyValid()) {
                 if (permit_uncompressed || pubkey.IsCompressed()) {
-                    return std::make_unique<ConstPubkeyProvider>(key_exp_index, pubkey);
+                    return std::make_unique<ConstPubkeyProvider>(key_exp_index, pubkey, false);
                 } else {
                     error = "Uncompressed keys are not allowed";
                     return nullptr;
@@ -952,7 +952,7 @@ std::unique_ptr<PubkeyProvider> ParsePubkeyInner(uint32_t key_exp_index, const S
             if (permit_uncompressed || key.IsCompressed()) {
                 CPubKey pubkey = key.GetPubKey();
                 out.keys.emplace(pubkey.GetID(), key);
-                return std::make_unique<ConstPubkeyProvider>(key_exp_index, pubkey);
+                return std::make_unique<ConstPubkeyProvider>(key_exp_index, pubkey, ctx == ParseScriptContext::P2TR);
             } else {
                 error = "Uncompressed keys are not allowed";
                 return nullptr;
@@ -1221,7 +1221,7 @@ std::unique_ptr<DescriptorImpl> ParseScript(uint32_t& key_exp_index, Span<const
 
 std::unique_ptr<PubkeyProvider> InferPubkey(const CPubKey& pubkey, ParseScriptContext, const SigningProvider& provider)
 {
-    std::unique_ptr<PubkeyProvider> key_provider = std::make_unique<ConstPubkeyProvider>(0, pubkey);
+    std::unique_ptr<PubkeyProvider> key_provider = std::make_unique<ConstPubkeyProvider>(0, pubkey, false);
     KeyOriginInfo info;
     if (provider.GetKeyOrigin(pubkey.GetID(), info)) {
         return std::make_unique<OriginPubkeyProvider>(0, std::move(info), std::move(key_provider));
@@ -1229,18 +1229,42 @@ std::unique_ptr<PubkeyProvider> InferPubkey(const CPubKey& pubkey, ParseScriptCo
     return key_provider;
 }
 
+std::unique_ptr<PubkeyProvider> InferXOnlyPubkey(const XOnlyPubKey& xkey, ParseScriptContext ctx, const SigningProvider& provider)
+{
+    unsigned char full_key[CPubKey::COMPRESSED_SIZE] = {0x02};
+    std::copy(xkey.begin(), xkey.end(), full_key + 1);
+    CPubKey pubkey(full_key);
+    std::unique_ptr<PubkeyProvider> key_provider = std::make_unique<ConstPubkeyProvider>(0, pubkey, true);
+    KeyOriginInfo info;
+    if (provider.GetKeyOrigin(pubkey.GetID(), info)) {
+        return std::make_unique<OriginPubkeyProvider>(0, std::move(info), std::move(key_provider));
+    } else {
+        full_key[0] = 0x03;
+        pubkey = CPubKey(full_key);
+        if (provider.GetKeyOrigin(pubkey.GetID(), info)) {
+            return std::make_unique<OriginPubkeyProvider>(0, std::move(info), std::move(key_provider));
+        }
+    }
+    return key_provider;
+}
+
 std::unique_ptr<DescriptorImpl> InferScript(const CScript& script, ParseScriptContext ctx, const SigningProvider& provider)
 {
+    if (ctx == ParseScriptContext::P2TR && script.size() == 34 && script[0] == 32 && script[33] == OP_CHECKSIG) {
+        XOnlyPubKey key{Span<const unsigned char>{script.data() + 1, script.data() + 33}};
+        return std::make_unique<PKDescriptor>(InferXOnlyPubkey(key, ctx, provider));
+    }
+
     std::vector<std::vector<unsigned char>> data;
     TxoutType txntype = Solver(script, data);
 
-    if (txntype == TxoutType::PUBKEY) {
+    if (txntype == TxoutType::PUBKEY && (ctx == ParseScriptContext::TOP || ctx == ParseScriptContext::P2SH || ctx == ParseScriptContext::P2WSH)) {
         CPubKey pubkey(data[0]);
         if (pubkey.IsValid()) {
             return std::make_unique<PKDescriptor>(InferPubkey(pubkey, ctx, provider));
         }
     }
-    if (txntype == TxoutType::PUBKEYHASH) {
+    if (txntype == TxoutType::PUBKEYHASH && (ctx == ParseScriptContext::TOP || ctx == ParseScriptContext::P2SH || ctx == ParseScriptContext::P2WSH)) {
         uint160 hash(data[0]);
         CKeyID keyid(hash);
         CPubKey pubkey;
@@ -1248,7 +1272,7 @@ std::unique_ptr<DescriptorImpl> InferScript(const CScript& script, ParseScriptCo
             return std::make_unique<PKHDescriptor>(InferPubkey(pubkey, ctx, provider));
         }
     }
-    if (txntype == TxoutType::WITNESS_V0_KEYHASH && ctx != ParseScriptContext::P2WSH) {
+    if (txntype == TxoutType::WITNESS_V0_KEYHASH && (ctx == ParseScriptContext::TOP || ctx == ParseScriptContext::P2SH)) {
         uint160 hash(data[0]);
         CKeyID keyid(hash);
         CPubKey pubkey;
@@ -1256,7 +1280,7 @@ std::unique_ptr<DescriptorImpl> InferScript(const CScript& script, ParseScriptCo
             return std::make_unique<WPKHDescriptor>(InferPubkey(pubkey, ctx, provider));
         }
     }
-    if (txntype == TxoutType::MULTISIG) {
+    if (txntype == TxoutType::MULTISIG && (ctx == ParseScriptContext::TOP || ctx == ParseScriptContext::P2SH || ctx == ParseScriptContext::P2WSH)) {
         std::vector<std::unique_ptr<PubkeyProvider>> providers;
         for (size_t i = 1; i + 1 < data.size(); ++i) {
             CPubKey pubkey(data[i]);
@@ -1273,7 +1297,7 @@ std::unique_ptr<DescriptorImpl> InferScript(const CScript& script, ParseScriptCo
             if (sub) return std::make_unique<SHDescriptor>(std::move(sub));
         }
     }
-    if (txntype == TxoutType::WITNESS_V0_SCRIPTHASH && ctx != ParseScriptContext::P2WSH) {
+    if (txntype == TxoutType::WITNESS_V0_SCRIPTHASH && (ctx == ParseScriptContext::TOP || ctx == ParseScriptContext::P2SH)) {
         CScriptID scriptid;
         CRIPEMD160().Write(data[0].data(), data[0].size()).Finalize(scriptid.begin());
         CScript subscript;
@@ -1282,6 +1306,40 @@ std::unique_ptr<DescriptorImpl> InferScript(const CScript& script, ParseScriptCo
             if (sub) return std::make_unique<WSHDescriptor>(std::move(sub));
         }
     }
+    if (txntype == TxoutType::WITNESS_V1_TAPROOT && ctx == ParseScriptContext::TOP) {
+        // Extract x-only pubkey from output.
+        XOnlyPubKey pubkey;
+        std::copy(data[0].begin(), data[0].end(), pubkey.begin());
+        // Request spending data.
+        TaprootSpendData tap;
+        if (provider.GetTaprootSpendData(pubkey, tap)) {
+            // If found, convert it back to tree form.
+            auto tree = InferTaprootTree(tap, pubkey);
+            if (tree) {
+                // If that works, try to infer subdescriptors for all leaves.
+                bool ok = true;
+                std::vector<std::unique_ptr<DescriptorImpl>> subscripts; //!< list of script subexpressions
+                std::vector<int> depths; //!< depth in the tree of each subexpression (same length subscripts)
+                for (const auto& [depth, script, leaf_ver] : *tree) {
+                    std::unique_ptr<DescriptorImpl> subdesc;
+                    if (leaf_ver == TAPROOT_LEAF_TAPSCRIPT) {
+                        subdesc = InferScript(script, ParseScriptContext::P2TR, provider);
+                    }
+                    if (!subdesc) {
+                        ok = false;
+                        break;
+                    } else {
+                        subscripts.push_back(std::move(subdesc));
+                        depths.push_back(depth);
+                    }
+                }
+                if (ok) {
+                    auto key = InferXOnlyPubkey(tap.internal_key, ParseScriptContext::P2TR, provider);
+                    return std::make_unique<TRDescriptor>(std::move(key), std::move(subscripts), std::move(depths));
+                }
+            }
+        }
+    }
 
     CTxDestination dest;
     if (ExtractDestination(script, dest)) {
diff --git a/src/script/standard.cpp b/src/script/standard.cpp
index 748f00dda5..b3dd5442fc 100644
--- a/src/script/standard.cpp
+++ b/src/script/standard.cpp
@@ -520,3 +520,138 @@ TaprootSpendData TaprootBuilder::GetSpendData() const
     }
     return spd;
 }
+
+std::optional<std::vector<std::tuple<int, CScript, int>>> InferTaprootTree(const TaprootSpendData& spenddata, const XOnlyPubKey& output)
+{
+    // Verify that the output matches the assumed Merkle root and internal key.
+    auto tweak = spenddata.internal_key.CreateTapTweak(spenddata.merkle_root.IsNull() ? nullptr : &spenddata.merkle_root);
+    if (!tweak || tweak->first != output) return std::nullopt;
+    // If the Merkle root is 0, the tree is empty, and we're done.
+    std::vector<std::tuple<int, CScript, int>> ret;
+    if (spenddata.merkle_root.IsNull()) return ret;
+
+    /** Data structure to represent the nodes of the tree we're going to be build. */
+    struct TreeNode {
+        /** Hash of this none, if known; 0 otherwise. */
+        uint256 hash;
+        /** The left and right subtrees (note that their order is irrelevant). */
+        std::unique_ptr<TreeNode> sub[2];
+        /** If this is known to be a leaf node, a pointer to the (script, leaf_ver) pair.
+         *  nullptr otherwise. */
+        const std::pair<CScript, int>* leaf = nullptr;
+        /** Whether or not this node has been explored (is known to be a leaf, or known to have children). */
+        bool explored = false;
+        /** Whether or not this node is an inner node (unknown until explored = true). */
+        bool inner;
+        /** Whether or not we have produced output for this subtree. */
+        bool done = false;
+    };
+
+    // Build tree from the provides branches.
+    TreeNode root;
+    root.hash = spenddata.merkle_root;
+    for (const auto& [key, control_blocks] : spenddata.scripts) {
+        const auto& [script, leaf_ver] = key;
+        for (const auto& control : control_blocks) {
+            // Skip script records with nonsensical leaf version.
+            if (leaf_ver < 0 || leaf_ver >= 0x100 || leaf_ver & 1) continue;
+            // Skip script records with invalid control block sizes.
+            if (control.size() < TAPROOT_CONTROL_BASE_SIZE || control.size() > TAPROOT_CONTROL_MAX_SIZE ||
+                ((control.size() - TAPROOT_CONTROL_BASE_SIZE) % TAPROOT_CONTROL_NODE_SIZE) != 0) continue;
+            // Skip script records that don't match the control block.
+            if ((control[0] & TAPROOT_LEAF_MASK) != leaf_ver) continue;
+            // Skip script records that don't match the provided Merkle root.
+            const uint256 leaf_hash = ComputeTapleafHash(leaf_ver, script);
+            const uint256 merkle_root = ComputeTaprootMerkleRoot(control, leaf_hash);
+            if (merkle_root != spenddata.merkle_root) continue;
+
+            TreeNode* node = &root;
+            size_t levels = (control.size() - TAPROOT_CONTROL_BASE_SIZE) / TAPROOT_CONTROL_NODE_SIZE;
+            for (size_t depth = 0; depth < levels; ++depth) {
+                // Can't descend into a node which we already know is a leaf.
+                if (node->explored && !node->inner) return std::nullopt;
+
+                // Extract partner hash from Merkle branch in control block.
+                uint256 hash;
+                std::copy(control.begin() + TAPROOT_CONTROL_BASE_SIZE + (levels - 1 - depth) * TAPROOT_CONTROL_NODE_SIZE,
+                          control.begin() + TAPROOT_CONTROL_BASE_SIZE + (levels - depth) * TAPROOT_CONTROL_NODE_SIZE,
+                          hash.begin());
+
+                if (node->sub[0]) {
+                    // Descend into the existing left or right branch.
+                    bool desc = false;
+                    for (int i = 0; i < 2; ++i) {
+                        if (node->sub[i]->hash == hash || (node->sub[i]->hash.IsNull() && node->sub[1-i]->hash != hash)) {
+                            node->sub[i]->hash = hash;
+                            node = &*node->sub[1-i];
+                            desc = true;
+                            break;
+                        }
+                    }
+                    if (!desc) return std::nullopt; // This probably requires a hash collision to hit.
+                } else {
+                    // We're in an unexplored node. Create subtrees and descend.
+                    node->explored = true;
+                    node->inner = true;
+                    node->sub[0] = std::make_unique<TreeNode>();
+                    node->sub[1] = std::make_unique<TreeNode>();
+                    node->sub[1]->hash = hash;
+                    node = &*node->sub[0];
+                }
+            }
+            // Cannot turn a known inner node into a leaf.
+            if (node->sub[0]) return std::nullopt;
+            node->explored = true;
+            node->inner = false;
+            node->leaf = &key;
+            node->hash = leaf_hash;
+        }
+    }
+
+    // Recursive processing to turn the tree into flattened output. Use an explicit stack here to avoid
+    // overflowing the call stack (the tree may be 128 levels deep).
+    std::vector<TreeNode*> stack{&root};
+    while (!stack.empty()) {
+        TreeNode& node = *stack.back();
+        if (!node.explored) {
+            // Unexplored node, which means the tree is incomplete.
+            return std::nullopt;
+        } else if (!node.inner) {
+            // Leaf node; produce output.
+            ret.emplace_back(stack.size() - 1, node.leaf->first, node.leaf->second);
+            node.done = true;
+            stack.pop_back();
+        } else if (node.sub[0]->done && !node.sub[1]->done && !node.sub[1]->explored && !node.sub[1]->hash.IsNull() &&
+                   (CHashWriter{HASHER_TAPBRANCH} << node.sub[1]->hash << node.sub[1]->hash).GetSHA256() == node.hash) {
+            // Whenever there are nodes with two identical subtrees under it, we run into a problem:
+            // the control blocks for the leaves underneath those will be identical as well, and thus
+            // they will all be matched to the same path in the tree. The result is that at the location
+            // where the duplicate occurred, the left child will contain a normal tree that can be explored
+            // and processed, but the right one will remain unexplored.
+            //
+            // This situation can be detected, by encountering an inner node with unexplored right subtree
+            // with known hash, and H_TapBranch(hash, hash) is equal to the parent node (this node)'s hash.
+            //
+            // To deal with this, simply process the left tree a second time (set its done flag to false;
+            // noting that the done flag of its children have already been set to false after processing
+            // those). To avoid ending up in an infinite loop, set the done flag of the right (unexplored)
+            // subtree to true.
+            node.sub[0]->done = false;
+            node.sub[1]->done = true;
+        } else if (node.sub[0]->done && node.sub[1]->done) {
+            // An internal node which we're finished with.
+            node.sub[0]->done = false;
+            node.sub[1]->done = false;
+            node.done = true;
+            stack.pop_back();
+        } else if (!node.sub[0]->done) {
+            // An internal node whose left branch hasn't been processed yet. Do so first.
+            stack.push_back(&*node.sub[0]);
+        } else if (!node.sub[1]->done) {
+            // An internal node whose right branch hasn't been processed yet. Do so first.
+            stack.push_back(&*node.sub[1]);
+        }
+    }
+
+    return ret;
+}
diff --git a/src/script/standard.h b/src/script/standard.h
index 285dd4c116..ac4e2f3276 100644
--- a/src/script/standard.h
+++ b/src/script/standard.h
@@ -327,4 +327,12 @@ public:
     TaprootSpendData GetSpendData() const;
 };
 
+/** Given a TaprootSpendData and the output key, reconstruct its script tree.
+ *
+ * If the output doesn't match the spenddata, or if the data in spenddata is incomplete,
+ * std::nullopt is returned. Otherwise, a vector of (depth, script, leaf_ver) tuples is
+ * returned, corresponding to a depth-first traversal of the script tree.
+ */
+std::optional<std::vector<std::tuple<int, CScript, int>>> InferTaprootTree(const TaprootSpendData& spenddata, const XOnlyPubKey& output);
+
 #endif // BITCOIN_SCRIPT_STANDARD_H