Load the descriptor cache from the wallet file

4 files changed, 64 insertions, 0 deletions

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index a6978d4aad..7142867081 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -1633,3 +1633,23 @@ void DescriptorScriptPubKeyMan::SetType(OutputType type, bool internal)
     this->m_address_type = type;
     this->m_internal = internal;
 }
+
+void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache)
+{
+    LOCK(cs_desc_man);
+    m_wallet_descriptor.cache = cache;
+    for (int32_t i = m_wallet_descriptor.range_start; i < m_wallet_descriptor.range_end; ++i) {
+        FlatSigningProvider out_keys;
+        std::vector<CScript> scripts_temp;
+        if (!m_wallet_descriptor.descriptor->ExpandFromCache(i, m_wallet_descriptor.cache, scripts_temp, out_keys)) {
+            throw std::runtime_error("Error: Unable to expand wallet descriptor from cache");
+        }
+        // Add all of the scriptPubKeys to the scriptPubKey set
+        for (const CScript& script : scripts_temp) {
+            if (m_map_script_pub_keys.count(script) != 0) {
+                throw std::runtime_error(strprintf("Error: Already loaded script at index %d as being at index %d", i, m_map_script_pub_keys[script]));
+            }
+            m_map_script_pub_keys[script] = i;
+        }
+    }
+}
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 81c89de985..34ba07967e 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -543,6 +543,8 @@ public:
     uint256 GetID() const override;
 
     void SetType(OutputType type, bool internal) override;
+
+    void SetCache(const DescriptorCache& cache);
 };
 
 #endif // BITCOIN_WALLET_SCRIPTPUBKEYMAN_H
diff --git a/src/wallet/walletdb.cpp b/src/wallet/walletdb.cpp
index c1e3694dcd..fc4783f932 100644
--- a/src/wallet/walletdb.cpp
+++ b/src/wallet/walletdb.cpp
@@ -44,6 +44,7 @@ const std::string SETTINGS{"settings"};
 const std::string TX{"tx"};
 const std::string VERSION{"version"};
 const std::string WALLETDESCRIPTOR{"walletdescriptor"};
+const std::string WALLETDESCRIPTORCACHE{"walletdescriptorcache"};
 const std::string WATCHMETA{"watchmeta"};
 const std::string WATCHS{"watchs"};
 } // namespace DBKeys
@@ -200,6 +201,7 @@ public:
     std::vector<uint256> vWalletUpgrade;
     std::map<OutputType, uint256> m_active_external_spks;
     std::map<OutputType, uint256> m_active_internal_spks;
+    std::map<uint256, DescriptorCache> m_descriptor_caches;
 
     CWalletScanState() {
     }
@@ -433,7 +435,38 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
             ssKey >> id;
             WalletDescriptor desc;
             ssValue >> desc;
+            if (wss.m_descriptor_caches.count(id) == 0) {
+                wss.m_descriptor_caches[id] = DescriptorCache();
+            }
             pwallet->LoadDescriptorScriptPubKeyMan(id, desc);
+        } else if (strType == DBKeys::WALLETDESCRIPTORCACHE) {
+            bool parent = true;
+            uint256 desc_id;
+            uint32_t key_exp_index;
+            uint32_t der_index;
+            ssKey >> desc_id;
+            ssKey >> key_exp_index;
+
+            // if the der_index exists, it's a derived xpub
+            try
+            {
+                ssKey >> der_index;
+                parent = false;
+            }
+            catch (...) {}
+
+            std::vector<unsigned char> ser_xpub(BIP32_EXTKEY_SIZE);
+            ssValue >> ser_xpub;
+            CExtPubKey xpub;
+            xpub.Decode(ser_xpub.data());
+            if (wss.m_descriptor_caches.count(desc_id)) {
+                wss.m_descriptor_caches[desc_id] = DescriptorCache();
+            }
+            if (parent) {
+                wss.m_descriptor_caches[desc_id].CacheParentExtPubKey(key_exp_index, xpub);
+            } else {
+                wss.m_descriptor_caches[desc_id].CacheDerivedExtPubKey(key_exp_index, der_index, xpub);
+            }
         } else if (strType != DBKeys::BESTBLOCK && strType != DBKeys::BESTBLOCK_NOMERKLE &&
                    strType != DBKeys::MINVERSION && strType != DBKeys::ACENTRY &&
                    strType != DBKeys::VERSION && strType != DBKeys::SETTINGS) {
@@ -535,6 +568,13 @@ DBErrors WalletBatch::LoadWallet(CWallet* pwallet)
         pwallet->SetActiveScriptPubKeyMan(spk_man_pair.second, spk_man_pair.first, /* internal */ true, /* memonly */ true);
     }
 
+    // Set the descriptor caches
+    for (auto desc_cache_pair : wss.m_descriptor_caches) {
+        auto spk_man = pwallet->GetScriptPubKeyMan(desc_cache_pair.first);
+        assert(spk_man);
+        ((DescriptorScriptPubKeyMan*)spk_man)->SetCache(desc_cache_pair.second);
+    }
+
     if (fNoncriticalErrors && result == DBErrors::LOAD_OK)
         result = DBErrors::NONCRITICAL_ERROR;
 
diff --git a/src/wallet/walletdb.h b/src/wallet/walletdb.h
index 4bd03da4a3..75fb12c9df 100644
--- a/src/wallet/walletdb.h
+++ b/src/wallet/walletdb.h
@@ -77,6 +77,8 @@ extern const std::string SETTINGS;
 extern const std::string TX;
 extern const std::string VERSION;
 extern const std::string WALLETDESCRIPTOR;
+extern const std::string WALLETDESCRIPTORCKEY;
+extern const std::string WALLETDESCRIPTORKEY;
 extern const std::string WATCHMETA;
 extern const std::string WATCHS;
 } // namespace DBKeys