1 files changed, 195 insertions, 3 deletions

diff --git a/src/script/standard.cpp b/src/script/standard.cpp
index a4b11cc0a9..b8349bb9ab 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,27 @@ 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) {
+        // Once P0083R3 is supported by all our targeted platforms,
+        // this loop body can be replaced with:
+        // scripts[key].merge(std::move(control_blocks));
+        auto& target = scripts[key];
+        for (auto& control_block: control_blocks) {
+            target.insert(std::move(control_block));
+        }
+    }
+}
+
 void TaprootBuilder::Insert(TaprootBuilder::NodeInfo&& node, int depth)
 {
     assert(depth >= 0 && (size_t)depth <= TAPROOT_CONTROL_MAX_NODE_COUNT);
@@ -435,13 +466,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 +496,168 @@ 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;
+}
+
+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 build. */
+    struct TreeNode {
+        /** Hash of this node, 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 provided 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;
+}