From 474ea3b927ddc67e64ae78e08c20c9264817e84d Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Mon, 24 Feb 2020 15:09:15 -0500
Subject: [PATCH] Introduce DescriptorCache struct which caches xpubs

---
 src/script/descriptor.cpp | 39 +++++++++++++++++++++++++++++++++++
 src/script/descriptor.h   | 43 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index b1d9a5bda7e..f71e3b901fe 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -1050,3 +1050,42 @@ std::unique_ptr<Descriptor> InferDescriptor(const CScript& script, const Signing
 {
     return InferScript(script, ParseScriptContext::TOP, provider);
 }
+
+void DescriptorCache::CacheParentExtPubKey(uint32_t key_exp_pos, const CExtPubKey& xpub)
+{
+    m_parent_xpubs[key_exp_pos] = xpub;
+}
+
+void DescriptorCache::CacheDerivedExtPubKey(uint32_t key_exp_pos, uint32_t der_index, const CExtPubKey& xpub)
+{
+    auto& xpubs = m_derived_xpubs[key_exp_pos];
+    xpubs[der_index] = xpub;
+}
+
+bool DescriptorCache::GetCachedParentExtPubKey(uint32_t key_exp_pos, CExtPubKey& xpub) const
+{
+    const auto& it = m_parent_xpubs.find(key_exp_pos);
+    if (it == m_parent_xpubs.end()) return false;
+    xpub = it->second;
+    return true;
+}
+
+bool DescriptorCache::GetCachedDerivedExtPubKey(uint32_t key_exp_pos, uint32_t der_index, CExtPubKey& xpub) const
+{
+    const auto& key_exp_it = m_derived_xpubs.find(key_exp_pos);
+    if (key_exp_it == m_derived_xpubs.end()) return false;
+    const auto& der_it = key_exp_it->second.find(der_index);
+    if (der_it == key_exp_it->second.end()) return false;
+    xpub = der_it->second;
+    return true;
+}
+
+const ExtPubKeyMap DescriptorCache::GetCachedParentExtPubKeys() const
+{
+    return m_parent_xpubs;
+}
+
+const std::unordered_map<uint32_t, ExtPubKeyMap> DescriptorCache::GetCachedDerivedExtPubKeys() const
+{
+    return m_derived_xpubs;
+}
diff --git a/src/script/descriptor.h b/src/script/descriptor.h
index 58b920c6812..5c686d68c10 100644
--- a/src/script/descriptor.h
+++ b/src/script/descriptor.h
@@ -13,6 +13,49 @@
 
 #include <vector>
 
+using ExtPubKeyMap = std::unordered_map<uint32_t, CExtPubKey>;
+
+/** Cache for single descriptor's derived extended pubkeys */
+class DescriptorCache {
+private:
+    /** Map key expression index -> map of (key derivation index -> xpub) */
+    std::unordered_map<uint32_t, ExtPubKeyMap> m_derived_xpubs;
+    /** Map key expression index -> parent xpub */
+    ExtPubKeyMap m_parent_xpubs;
+
+public:
+    /** Cache a parent xpub
+     *
+     * @param[in] key_exp_pos Position of the key expression within the descriptor
+     * @param[in] xpub The CExtPubKey to cache
+     */
+    void CacheParentExtPubKey(uint32_t key_exp_pos, const CExtPubKey& xpub);
+    /** Retrieve a cached parent xpub
+     *
+     * @param[in] key_exp_pos Position of the key expression within the descriptor
+     * @param[in] xpub The CExtPubKey to get from cache
+     */
+    bool GetCachedParentExtPubKey(uint32_t key_exp_pos, CExtPubKey& xpub) const;
+    /** Cache an xpub derived at an index
+     *
+     * @param[in] key_exp_pos Position of the key expression within the descriptor
+     * @param[in] der_index Derivation index of the xpub
+     * @param[in] xpub The CExtPubKey to cache
+     */
+    void CacheDerivedExtPubKey(uint32_t key_exp_pos, uint32_t der_index, const CExtPubKey& xpub);
+    /** Retrieve a cached xpub derived at an index
+     *
+     * @param[in] key_exp_pos Position of the key expression within the descriptor
+     * @param[in] der_index Derivation index of the xpub
+     * @param[in] xpub The CExtPubKey to get from cache
+     */
+    bool GetCachedDerivedExtPubKey(uint32_t key_exp_pos, uint32_t der_index, CExtPubKey& xpub) const;
+
+    /** Retrieve all cached parent xpubs */
+    const ExtPubKeyMap GetCachedParentExtPubKeys() const;
+    /** Retrieve all cached derived xpubs */
+    const std::unordered_map<uint32_t, ExtPubKeyMap> GetCachedDerivedExtPubKeys() const;
+};
 
 /** \brief Interface for parsed descriptor objects.
  *