From 90fcac365e1616779b40a69736428435df75fdf2 Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter@wuille.net>
Date: Mon, 1 Feb 2021 18:53:24 -0800
Subject: Add TaprootBuilder class

This class functions as a utility for building taproot outputs, from
internal key and script leaves.
---
 src/script/standard.h | 78 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)

(limited to 'src/script/standard.h')

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
-- 
cgit v1.2.3