From 3ae5b6af21cf45b3da5e341e84f50e0717eaf589 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Mon, 19 Jul 2021 16:01:12 -0400
Subject: Store TaprootBuilder in SigningProviders instead of TaprootSpendData

TaprootSpendData can be gotten from TaprootBuilder, however for PSBT, we
also need TaprootBuilders directly (for the outputs). So we store the
TaprootBuilder in the FlatSigningProvider and when the TaprootSpendData
is needed, we generate it on the fly using the stored builder.
---
 src/script/descriptor.cpp      |  2 +-
 src/script/signingprovider.cpp | 21 ++++++++++++++++-----
 src/script/signingprovider.h   |  5 ++++-
 3 files changed, 21 insertions(+), 7 deletions(-)

(limited to 'src/script')

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index cece0b60ce..ca0170c84b 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -882,7 +882,7 @@ protected:
         if (!xpk.IsFullyValid()) return {};
         builder.Finalize(xpk);
         WitnessV1Taproot output = builder.GetOutput();
-        out.tr_spenddata[output].Merge(builder.GetSpendData());
+        out.tr_trees[output] = builder;
         out.pubkeys.emplace(keys[0].GetID(), keys[0]);
         return Vector(GetScriptForDestination(output));
     }
diff --git a/src/script/signingprovider.cpp b/src/script/signingprovider.cpp
index 552934e0eb..c624a17628 100644
--- a/src/script/signingprovider.cpp
+++ b/src/script/signingprovider.cpp
@@ -48,6 +48,10 @@ bool HidingSigningProvider::GetTaprootSpendData(const XOnlyPubKey& output_key, T
 {
     return m_provider->GetTaprootSpendData(output_key, spenddata);
 }
+bool HidingSigningProvider::GetTaprootBuilder(const XOnlyPubKey& output_key, TaprootBuilder& builder) const
+{
+    return m_provider->GetTaprootBuilder(output_key, builder);
+}
 
 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); }
@@ -61,7 +65,16 @@ bool FlatSigningProvider::GetKeyOrigin(const CKeyID& keyid, KeyOriginInfo& info)
 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);
+    TaprootBuilder builder;
+    if (LookupHelper(tr_trees, output_key, builder)) {
+        spenddata = builder.GetSpendData();
+        return true;
+    }
+    return false;
+}
+bool FlatSigningProvider::GetTaprootBuilder(const XOnlyPubKey& output_key, TaprootBuilder& builder) const
+{
+    return LookupHelper(tr_trees, output_key, builder);
 }
 
 FlatSigningProvider Merge(const FlatSigningProvider& a, const FlatSigningProvider& b)
@@ -75,10 +88,8 @@ 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);
-    }
+    ret.tr_trees = a.tr_trees;
+    ret.tr_trees.insert(b.tr_trees.begin(), b.tr_trees.end());
     return ret;
 }
 
diff --git a/src/script/signingprovider.h b/src/script/signingprovider.h
index f1bded1a8c..792cc903f2 100644
--- a/src/script/signingprovider.h
+++ b/src/script/signingprovider.h
@@ -25,6 +25,7 @@ public:
     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; }
+    virtual bool GetTaprootBuilder(const XOnlyPubKey& output_key, TaprootBuilder& builder) const { return false; }
 
     bool GetKeyByXOnly(const XOnlyPubKey& pubkey, CKey& key) const
     {
@@ -67,6 +68,7 @@ public:
     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;
+    bool GetTaprootBuilder(const XOnlyPubKey& output_key, TaprootBuilder& builder) const override;
 };
 
 struct FlatSigningProvider final : public SigningProvider
@@ -75,13 +77,14 @@ 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. */
+    std::map<XOnlyPubKey, TaprootBuilder> tr_trees; /** Map from output key to Taproot tree (which can then make the TaprootSpendData */
 
     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;
+    bool GetTaprootBuilder(const XOnlyPubKey& output_key, TaprootBuilder& builder) const override;
 };
 
 FlatSigningProvider Merge(const FlatSigningProvider& a, const FlatSigningProvider& b);
-- 
cgit v1.2.3