Add descriptor expansion cache

3 files changed, 75 insertions, 30 deletions

diff --git a/src/script/descriptor.cpp b/src/script/descriptor.cpp
index adb5d20edb..a702be5b78 100644
--- a/src/script/descriptor.cpp
+++ b/src/script/descriptor.cpp
@@ -40,8 +40,8 @@ struct PubkeyProvider
 {
     virtual ~PubkeyProvider() = default;
 
-    /** Derive a public key. */
-    virtual bool GetPubKey(int pos, const SigningProvider& arg, CPubKey& key, KeyOriginInfo& info) const = 0;
+    /** Derive a public key. If key==nullptr, only info is desired. */
+    virtual bool GetPubKey(int pos, const SigningProvider& arg, CPubKey* key, KeyOriginInfo& info) const = 0;
 
     /** Whether this represent multiple public keys at different positions. */
     virtual bool IsRange() const = 0;
@@ -68,7 +68,7 @@ class OriginPubkeyProvider final : public PubkeyProvider
 
 public:
     OriginPubkeyProvider(KeyOriginInfo info, std::unique_ptr<PubkeyProvider> provider) : m_origin(std::move(info)), m_provider(std::move(provider)) {}
-    bool GetPubKey(int pos, const SigningProvider& arg, CPubKey& key, KeyOriginInfo& info) const override
+    bool GetPubKey(int pos, const SigningProvider& arg, CPubKey* key, KeyOriginInfo& info) const override
     {
         if (!m_provider->GetPubKey(pos, arg, key, info)) return false;
         std::copy(std::begin(m_origin.fingerprint), std::end(m_origin.fingerprint), info.fingerprint);
@@ -94,9 +94,9 @@ class ConstPubkeyProvider final : public PubkeyProvider
 
 public:
     ConstPubkeyProvider(const CPubKey& pubkey) : m_pubkey(pubkey) {}
-    bool GetPubKey(int pos, const SigningProvider& arg, CPubKey& key, KeyOriginInfo& info) const override
+    bool GetPubKey(int pos, const SigningProvider& arg, CPubKey* key, KeyOriginInfo& info) const override
     {
-        key = m_pubkey;
+        if (key) *key = m_pubkey;
         info.path.clear();
         CKeyID keyid = m_pubkey.GetID();
         std::copy(keyid.begin(), keyid.begin() + sizeof(info.fingerprint), info.fingerprint);
@@ -152,26 +152,28 @@ public:
     BIP32PubkeyProvider(const CExtPubKey& extkey, KeyPath path, DeriveType derive) : m_extkey(extkey), m_path(std::move(path)), m_derive(derive) {}
     bool IsRange() const override { return m_derive != DeriveType::NO; }
     size_t GetSize() const override { return 33; }
-    bool GetPubKey(int pos, const SigningProvider& arg, CPubKey& key, KeyOriginInfo& info) const override
+    bool GetPubKey(int pos, const SigningProvider& arg, CPubKey* key, KeyOriginInfo& info) const override
     {
-        if (IsHardened()) {
-            CExtKey extkey;
-            if (!GetExtKey(arg, extkey)) return false;
-            for (auto entry : m_path) {
-                extkey.Derive(extkey, entry);
-            }
-            if (m_derive == DeriveType::UNHARDENED) extkey.Derive(extkey, pos);
-            if (m_derive == DeriveType::HARDENED) extkey.Derive(extkey, pos | 0x80000000UL);
-            key = extkey.Neuter().pubkey;
-        } else {
-            // TODO: optimize by caching
-            CExtPubKey extkey = m_extkey;
-            for (auto entry : m_path) {
-                extkey.Derive(extkey, entry);
+        if (key) {
+            if (IsHardened()) {
+                CExtKey extkey;
+                if (!GetExtKey(arg, extkey)) return false;
+                for (auto entry : m_path) {
+                    extkey.Derive(extkey, entry);
+                }
+                if (m_derive == DeriveType::UNHARDENED) extkey.Derive(extkey, pos);
+                if (m_derive == DeriveType::HARDENED) extkey.Derive(extkey, pos | 0x80000000UL);
+                *key = extkey.Neuter().pubkey;
+            } else {
+                // TODO: optimize by caching
+                CExtPubKey extkey = m_extkey;
+                for (auto entry : m_path) {
+                    extkey.Derive(extkey, entry);
+                }
+                if (m_derive == DeriveType::UNHARDENED) extkey.Derive(extkey, pos);
+                assert(m_derive != DeriveType::HARDENED);
+                *key = extkey.pubkey;
             }
-            if (m_derive == DeriveType::UNHARDENED) extkey.Derive(extkey, pos);
-            assert(m_derive != DeriveType::HARDENED);
-            key = extkey.pubkey;
         }
         CKeyID keyid = m_extkey.pubkey.GetID();
         std::copy(keyid.begin(), keyid.begin() + sizeof(info.fingerprint), info.fingerprint);
@@ -285,7 +287,7 @@ public:
 
     bool ToPrivateString(const SigningProvider& arg, std::string& out) const override final { return ToStringHelper(&arg, out, true); }
 
-    bool Expand(int pos, const SigningProvider& arg, std::vector<CScript>& output_scripts, FlatSigningProvider& out) const final
+    bool ExpandHelper(int pos, const SigningProvider& arg, Span<const unsigned char>* cache_read, std::vector<CScript>& output_scripts, FlatSigningProvider& out, std::vector<unsigned char>* cache_write) const
     {
         std::vector<std::pair<CPubKey, KeyOriginInfo>> entries;
         entries.reserve(m_pubkey_args.size());
@@ -293,12 +295,25 @@ public:
         // Construct temporary data in `entries` and `subscripts`, to avoid producing output in case of failure.
         for (const auto& p : m_pubkey_args) {
             entries.emplace_back();
-            if (!p->GetPubKey(pos, arg, entries.back().first, entries.back().second)) return false;
+            if (!p->GetPubKey(pos, arg, cache_read ? nullptr : &entries.back().first, entries.back().second)) return false;
+            if (cache_read) {
+                // Cached expanded public key exists, use it.
+                if (cache_read->size() == 0) return false;
+                bool compressed = ((*cache_read)[0] == 0x02 || (*cache_read)[0] == 0x03) && cache_read->size() >= 33;
+                bool uncompressed = ((*cache_read)[0] == 0x04) && cache_read->size() >= 65;
+                if (!(compressed || uncompressed)) return false;
+                CPubKey pubkey(cache_read->begin(), cache_read->begin() + (compressed ? 33 : 65));
+                entries.back().first = pubkey;
+                *cache_read = cache_read->subspan(compressed ? 33 : 65);
+            }
+            if (cache_write) {
+                cache_write->insert(cache_write->end(), entries.back().first.begin(), entries.back().first.end());
+            }
         }
         std::vector<CScript> subscripts;
         if (m_script_arg) {
             FlatSigningProvider subprovider;
-            if (!m_script_arg->Expand(pos, arg, subscripts, subprovider)) return false;
+            if (!m_script_arg->ExpandHelper(pos, arg, cache_read, subscripts, subprovider, cache_write)) return false;
             out = Merge(out, subprovider);
         }
 
@@ -322,6 +337,17 @@ public:
         }
         return true;
     }
+
+    bool Expand(int pos, const SigningProvider& provider, std::vector<CScript>& output_scripts, FlatSigningProvider& out, std::vector<unsigned char>* cache = nullptr) const final
+    {
+        return ExpandHelper(pos, provider, nullptr, output_scripts, out, cache);
+    }
+
+    bool ExpandFromCache(int pos, const std::vector<unsigned char>& cache, std::vector<CScript>& output_scripts, FlatSigningProvider& out) const final
+    {
+        Span<const unsigned char> span = MakeSpan(cache);
+        return ExpandHelper(pos, DUMMY_SIGNING_PROVIDER, &span, output_scripts, out, nullptr) && span.size() == 0;
+    }
 };
 
 /** Construct a vector with one element, which is moved into it. */
diff --git a/src/script/descriptor.h b/src/script/descriptor.h
index 0111972f85..44f0efca03 100644
--- a/src/script/descriptor.h
+++ b/src/script/descriptor.h
@@ -48,8 +48,18 @@ struct Descriptor {
      * provider: the provider to query for private keys in case of hardened derivation.
      * output_script: the expanded scriptPubKeys will be put here.
      * out: scripts and public keys necessary for solving the expanded scriptPubKeys will be put here (may be equal to provider).
+     * cache: vector which will be overwritten with cache data necessary to-evaluate the descriptor at this point without access to private keys.
      */
-    virtual bool Expand(int pos, const SigningProvider& provider, std::vector<CScript>& output_scripts, FlatSigningProvider& out) const = 0;
+    virtual bool Expand(int pos, const SigningProvider& provider, std::vector<CScript>& output_scripts, FlatSigningProvider& out, std::vector<unsigned char>* cache = nullptr) const = 0;
+
+    /** Expand a descriptor at a specified position using cached expansion data.
+     *
+     * pos: the position at which to expand the descriptor. If IsRange() is false, this is ignored.
+     * cache: vector from which cached expansion data will be read.
+     * output_script: the expanded scriptPubKeys will be put here.
+     * out: scripts and public keys necessary for solving the expanded scriptPubKeys will be put here (may be equal to provider).
+     */
+    virtual bool ExpandFromCache(int pos, const std::vector<unsigned char>& cache, std::vector<CScript>& output_scripts, FlatSigningProvider& out) const = 0;
 };
 
 /** Parse a descriptor string. Included private keys are put in out. Returns nullptr if parsing fails. */
diff --git a/src/test/descriptor_tests.cpp b/src/test/descriptor_tests.cpp
index 0e98f5a826..1a35b5ea71 100644
--- a/src/test/descriptor_tests.cpp
+++ b/src/test/descriptor_tests.cpp
@@ -88,9 +88,18 @@ void Check(const std::string& prv, const std::string& pub, int flags, const std:
         const auto& ref = scripts[(flags & RANGE) ? i : 0];
         for (int t = 0; t < 2; ++t) {
             const FlatSigningProvider& key_provider = (flags & HARDENED) ? keys_priv : keys_pub;
-            FlatSigningProvider script_provider;
-            std::vector<CScript> spks;
-            BOOST_CHECK((t ? parse_priv : parse_pub)->Expand(i, key_provider, spks, script_provider));
+            FlatSigningProvider script_provider, script_provider_cached;
+            std::vector<CScript> spks, spks_cached;
+            std::vector<unsigned char> cache;
+            BOOST_CHECK((t ? parse_priv : parse_pub)->Expand(i, key_provider, spks, script_provider, &cache));
+
+            // Try to expand again using cached data, and compare.
+            BOOST_CHECK(parse_pub->ExpandFromCache(i, cache, spks_cached, script_provider_cached));
+            BOOST_CHECK(spks == spks_cached);
+            BOOST_CHECK(script_provider.pubkeys == script_provider_cached.pubkeys);
+            BOOST_CHECK(script_provider.scripts == script_provider_cached.scripts);
+            BOOST_CHECK(script_provider.origins == script_provider_cached.origins);
+
             BOOST_CHECK_EQUAL(spks.size(), ref.size());
             for (size_t n = 0; n < spks.size(); ++n) {
                 BOOST_CHECK_EQUAL(ref[n], HexStr(spks[n].begin(), spks[n].end()));