From ac7747585fb0629be502a089c9c9be876bd7107d Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Mon, 19 Jul 2021 16:02:36 -0400
Subject: Fill PSBT Taproot output data to/from SignatureData

---
 src/psbt.cpp        | 18 ++++++++++++++++++
 src/script/sign.cpp |  4 ++++
 src/script/sign.h   |  1 +
 3 files changed, 23 insertions(+)

(limited to 'src')

diff --git a/src/psbt.cpp b/src/psbt.cpp
index ba32d7acae..e4bac4aa63 100644
--- a/src/psbt.cpp
+++ b/src/psbt.cpp
@@ -213,6 +213,15 @@ void PSBTOutput::FillSignatureData(SignatureData& sigdata) const
     for (const auto& key_pair : hd_keypaths) {
         sigdata.misc_pubkeys.emplace(key_pair.first.GetID(), key_pair);
     }
+    if (m_tap_tree.has_value() && m_tap_internal_key.IsFullyValid()) {
+        TaprootSpendData spenddata = m_tap_tree->GetSpendData();
+
+        sigdata.tr_spenddata.internal_key = m_tap_internal_key;
+        sigdata.tr_spenddata.Merge(spenddata);
+    }
+    for (const auto& [pubkey, leaf_origin] : m_tap_bip32_paths) {
+        sigdata.taproot_misc_pubkeys.emplace(pubkey, leaf_origin);
+    }
 }
 
 void PSBTOutput::FromSignatureData(const SignatureData& sigdata)
@@ -226,6 +235,15 @@ void PSBTOutput::FromSignatureData(const SignatureData& sigdata)
     for (const auto& entry : sigdata.misc_pubkeys) {
         hd_keypaths.emplace(entry.second);
     }
+    if (!sigdata.tr_spenddata.internal_key.IsNull()) {
+        m_tap_internal_key = sigdata.tr_spenddata.internal_key;
+    }
+    if (sigdata.tr_builder.has_value()) {
+        m_tap_tree = sigdata.tr_builder;
+    }
+    for (const auto& [pubkey, leaf_origin] : sigdata.taproot_misc_pubkeys) {
+        m_tap_bip32_paths.emplace(pubkey, leaf_origin);
+    }
 }
 
 bool PSBTOutput::IsNull() const
diff --git a/src/script/sign.cpp b/src/script/sign.cpp
index e46d9be7d4..5cd773f6ee 100644
--- a/src/script/sign.cpp
+++ b/src/script/sign.cpp
@@ -216,11 +216,15 @@ static bool SignTaprootScript(const SigningProvider& provider, const BaseSignatu
 static bool SignTaproot(const SigningProvider& provider, const BaseSignatureCreator& creator, const WitnessV1Taproot& output, SignatureData& sigdata, std::vector<valtype>& result)
 {
     TaprootSpendData spenddata;
+    TaprootBuilder builder;
 
     // Gather information about this output.
     if (provider.GetTaprootSpendData(output, spenddata)) {
         sigdata.tr_spenddata.Merge(spenddata);
     }
+    if (provider.GetTaprootBuilder(output, builder)) {
+        sigdata.tr_builder = builder;
+    }
 
     // Try key path spending.
     {
diff --git a/src/script/sign.h b/src/script/sign.h
index cb3c229298..5e58272154 100644
--- a/src/script/sign.h
+++ b/src/script/sign.h
@@ -70,6 +70,7 @@ struct SignatureData {
     CScript witness_script; ///< The witnessScript (if any) for the input. witnessScripts are used in P2WSH outputs.
     CScriptWitness scriptWitness; ///< The scriptWitness of an input. Contains complete signatures or the traditional partial signatures format. scriptWitness is part of a transaction input per BIP 144.
     TaprootSpendData tr_spenddata; ///< Taproot spending data.
+    std::optional<TaprootBuilder> tr_builder; ///< Taproot tree used to build tr_spenddata.
     std::map<CKeyID, SigPair> signatures; ///< BIP 174 style partial signatures for the input. May contain all signatures necessary for producing a final scriptSig or scriptWitness.
     std::map<CKeyID, std::pair<CPubKey, KeyOriginInfo>> misc_pubkeys;
     std::vector<unsigned char> taproot_key_path_sig; /// Schnorr signature for key path spending
-- 
cgit v1.2.3