Add caching of incremental witnesses for spendable notes

src/serialize.h

@@ -12,6 +12,7 @@
 #include <assert.h>
 #include <ios>
 #include <limits>
+#include <list>
 #include <map>
 #include <set>
 #include <stdint.h>
@@ -544,6 +545,13 @@ template<typename K, typename Pred, typename A> unsigned int GetSerializeSize(co
 template<typename Stream, typename K, typename Pred, typename A> void Serialize(Stream& os, const std::set<K, Pred, A>& m, int nType, int nVersion);
 template<typename Stream, typename K, typename Pred, typename A> void Unserialize(Stream& is, std::set<K, Pred, A>& m, int nType, int nVersion);
 
+/**
+ * list
+ */
+template<typename T, typename A> unsigned int GetSerializeSize(const std::list<T, A>& m, int nType, int nVersion);
+template<typename Stream, typename T, typename A> void Serialize(Stream& os, const std::list<T, A>& m, int nType, int nVersion);
+template<typename Stream, typename T, typename A> void Unserialize(Stream& is, std::list<T, A>& m, int nType, int nVersion);
+
 
 
 
@@ -890,6 +898,42 @@ void Unserialize(Stream& is, std::set<K, Pred, A>& m, int nType, int nVersion)
 
 
 
+/**
+ * list
+ */
+template<typename T, typename A>
+unsigned int GetSerializeSize(const std::list<T, A>& l, int nType, int nVersion)
+{
+    unsigned int nSize = GetSizeOfCompactSize(l.size());
+    for (typename std::list<T, A>::const_iterator it = l.begin(); it != l.end(); ++it)
+        nSize += GetSerializeSize((*it), nType, nVersion);
+    return nSize;
+}
+
+template<typename Stream, typename T, typename A>
+void Serialize(Stream& os, const std::list<T, A>& l, int nType, int nVersion)
+{
+    WriteCompactSize(os, l.size());
+    for (typename std::list<T, A>::const_iterator it = l.begin(); it != l.end(); ++it)
+        Serialize(os, (*it), nType, nVersion);
+}
+
+template<typename Stream, typename T, typename A>
+void Unserialize(Stream& is, std::list<T, A>& l, int nType, int nVersion)
+{
+    l.clear();
+    unsigned int nSize = ReadCompactSize(is);
+    typename std::list<T, A>::iterator it = l.begin();
+    for (unsigned int i = 0; i < nSize; i++)
+    {
+        T item;
+        Unserialize(is, item, nType, nVersion);
+        l.push_back(item);
+    }
+}
+
+
+
 /**
  * Support for ADD_SERIALIZE_METHODS and READWRITE macro
  */

src/wallet/gtest/test_wallet.cpp

@@ -1,3 +1,4 @@
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <sodium.h>
 
@@ -10,8 +11,32 @@
 #include "zcash/Note.hpp"
 #include "zcash/NoteEncryption.hpp"
 
+using ::testing::_;
+using ::testing::Return;
+
 ZCJoinSplit* params = ZCJoinSplit::Unopened();
 
+class MockCCoinsViewCache : public CCoinsViewCache {
+public:
+    MockCCoinsViewCache() : CCoinsViewCache(NULL) { };
+
+    MOCK_CONST_METHOD2(GetAnchorAt, bool(const uint256 &rt, ZCIncrementalMerkleTree &tree));
+};
+
+class TestWallet : public CWallet {
+public:
+    TestWallet() : CWallet() { }
+
+    void IncrementNoteWitnesses(const CBlockIndex* pindex,
+                                const CBlock* pblock,
+                                const CCoinsViewCache* pcoins) {
+        CWallet::IncrementNoteWitnesses(pindex, pblock, pcoins);
+    }
+    void DecrementNoteWitnesses() {
+        CWallet::DecrementNoteWitnesses();
+    }
+};
+
 CWalletTx GetValidReceive(const libzcash::SpendingKey& sk, CAmount value, bool randomInputs) {
     CMutableTransaction mtx;
     mtx.nVersion = 2; // Enable JoinSplits
@@ -272,3 +297,132 @@ TEST(wallet_tests, navigate_from_nullifier_to_note) {
     EXPECT_EQ(0, wallet.mapNullifiers[nullifier].js);
     EXPECT_EQ(1, wallet.mapNullifiers[nullifier].n);
 }
+
+TEST(wallet_tests, cached_witnesses_empty_chain) {
+    TestWallet wallet;
+
+    auto sk = libzcash::SpendingKey::random();
+    wallet.AddSpendingKey(sk);
+
+    auto wtx = GetValidReceive(sk, 10, true);
+    auto note = GetNote(sk, wtx, 0, 1);
+    auto nullifier = note.nullifier(sk);
+
+    mapNoteData_t noteData;
+    JSOutPoint jsoutpt {wtx.GetTxid(), 0, 1};
+    CNoteData nd {sk.address(), nullifier};
+    noteData[jsoutpt] = nd;
+    wtx.SetNoteData(noteData);
+
+    std::vector<JSOutPoint> notes {jsoutpt};
+    std::vector<boost::optional<ZCIncrementalWitness>> witnesses;
+    uint256 anchor;
+
+    wallet.GetNoteWitnesses(notes, witnesses, anchor);
+    EXPECT_FALSE((bool) witnesses[0]);
+
+    wallet.AddToWallet(wtx, true, NULL);
+    witnesses.clear();
+    wallet.GetNoteWitnesses(notes, witnesses, anchor);
+    EXPECT_FALSE((bool) witnesses[0]);
+
+    CBlock block;
+    block.vtx.push_back(wtx);
+    MockCCoinsViewCache coins;
+    // Empty chain, so we shouldn't try to fetch an anchor
+    EXPECT_CALL(coins, GetAnchorAt(_, _)).Times(0);
+    wallet.IncrementNoteWitnesses(NULL, &block, &coins);
+    witnesses.clear();
+    wallet.GetNoteWitnesses(notes, witnesses, anchor);
+    EXPECT_TRUE((bool) witnesses[0]);
+
+    wallet.DecrementNoteWitnesses();
+    witnesses.clear();
+    wallet.GetNoteWitnesses(notes, witnesses, anchor);
+    EXPECT_FALSE((bool) witnesses[0]);
+}
+
+TEST(wallet_tests, cached_witnesses_chain_tip) {
+    TestWallet wallet;
+    MockCCoinsViewCache coins;
+    uint256 anchor1;
+    CBlock block1;
+
+    auto sk = libzcash::SpendingKey::random();
+    wallet.AddSpendingKey(sk);
+
+    {
+        // First transaction (case tested in _empty_chain)
+        auto wtx = GetValidReceive(sk, 10, true);
+        auto note = GetNote(sk, wtx, 0, 1);
+        auto nullifier = note.nullifier(sk);
+
+        mapNoteData_t noteData;
+        JSOutPoint jsoutpt {wtx.GetTxid(), 0, 1};
+        CNoteData nd {sk.address(), nullifier};
+        noteData[jsoutpt] = nd;
+        wtx.SetNoteData(noteData);
+        wallet.AddToWallet(wtx, true, NULL);
+
+        std::vector<JSOutPoint> notes {jsoutpt};
+        std::vector<boost::optional<ZCIncrementalWitness>> witnesses;
+
+        // First block (case tested in _empty_chain)
+        block1.vtx.push_back(wtx);
+        EXPECT_CALL(coins, GetAnchorAt(_, _))
+            .Times(0);
+        wallet.IncrementNoteWitnesses(NULL, &block1, &coins);
+        // Called to fetch anchor
+        wallet.GetNoteWitnesses(notes, witnesses, anchor1);
+    }
+
+    {
+        // Second transaction
+        auto wtx = GetValidReceive(sk, 50, true);
+        auto note = GetNote(sk, wtx, 0, 1);
+        auto nullifier = note.nullifier(sk);
+
+        mapNoteData_t noteData;
+        JSOutPoint jsoutpt {wtx.GetTxid(), 0, 1};
+        CNoteData nd {sk.address(), nullifier};
+        noteData[jsoutpt] = nd;
+        wtx.SetNoteData(noteData);
+        wallet.AddToWallet(wtx, true, NULL);
+
+        std::vector<JSOutPoint> notes {jsoutpt};
+        std::vector<boost::optional<ZCIncrementalWitness>> witnesses;
+        uint256 anchor2;
+
+        wallet.GetNoteWitnesses(notes, witnesses, anchor2);
+        EXPECT_FALSE((bool) witnesses[0]);
+
+        // Second block
+        CBlock block2;
+        block2.hashPrevBlock = block1.GetHash();
+        block2.vtx.push_back(wtx);
+        EXPECT_CALL(coins, GetAnchorAt(anchor1, _))
+            .Times(2)
+            .WillRepeatedly(Return(true));
+        wallet.IncrementNoteWitnesses(NULL, &block2, &coins);
+        witnesses.clear();
+        wallet.GetNoteWitnesses(notes, witnesses, anchor2);
+        EXPECT_TRUE((bool) witnesses[0]);
+        EXPECT_NE(anchor1, anchor2);
+
+        // Decrementing should give us the previous anchor
+        uint256 anchor3;
+        wallet.DecrementNoteWitnesses();
+        witnesses.clear();
+        wallet.GetNoteWitnesses(notes, witnesses, anchor3);
+        EXPECT_FALSE((bool) witnesses[0]);
+        EXPECT_EQ(anchor1, anchor3);
+
+        // Re-incrementing with the same block should give the same result
+        uint256 anchor4;
+        wallet.IncrementNoteWitnesses(NULL, &block2, &coins);
+        witnesses.clear();
+        wallet.GetNoteWitnesses(notes, witnesses, anchor4);
+        EXPECT_TRUE((bool) witnesses[0]);
+        EXPECT_EQ(anchor2, anchor4);
+    }
+}

src/wallet/wallet.cpp

@@ -583,6 +583,110 @@ void CWallet::AddToSpends(const uint256& wtxid)
     }
 }
 
+void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
+                                     const CBlock* pblockIn,
+                                     const CCoinsViewCache* pcoins)
+{
+    {
+        LOCK(cs_wallet);
+        for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
+            for (mapNoteData_t::value_type& item : wtxItem.second.mapNoteData) {
+                CNoteData* nd = &(item.second);
+                // Copy the witness for the previous block if we have one
+                if (nd->witnesses.size() > 0) {
+                    nd->witnesses.push_front(nd->witnesses.front());
+                }
+                if (nd->witnesses.size() > WITNESS_CACHE_SIZE) {
+                    nd->witnesses.pop_back();
+                }
+            }
+        }
+
+        const CBlock* pblock {pblockIn};
+        CBlock block;
+        if (!pblock) {
+            ReadBlockFromDisk(block, pindex);
+            pblock = &block;
+        }
+
+        ZCIncrementalMerkleTree tree;
+        bool treeInitialised = false;
+        for (const CTransaction& tx : pblock->vtx) {
+            if (!treeInitialised && tx.vjoinsplit.size() > 0) {
+                LOCK(cs_main);
+                // vAnchorCache will only be empty at the beginning
+                if (vAnchorCache.size() && !pcoins->GetAnchorAt(vAnchorCache.front(), tree)) {
+                    // This should not happen, because IncrementNoteWitnesses()
+                    // is only called when the chain tip updates, and the
+                    // anchors for the JoinSplits in that block should still be
+                    // cached.
+                    // TODO: Calculate the anchor from scratch?
+                    throw std::runtime_error("CWallet::IncrementNoteWitnesses(): anchor not cached");
+                }
+                treeInitialised = true;
+            }
+
+            auto hash = tx.GetTxid();
+            bool txIsOurs = mapWallet.count(hash);
+            for (size_t i = 0; i < tx.vjoinsplit.size(); i++) {
+                const JSDescription& jsdesc = tx.vjoinsplit[i];
+                for (uint8_t j = 0; j < jsdesc.commitments.size(); j++) {
+                    const uint256& note_commitment = jsdesc.commitments[j];
+                    tree.append(note_commitment);
+
+                    // Increment existing witnesses
+                    for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
+                        for (mapNoteData_t::value_type& item : wtxItem.second.mapNoteData) {
+                            CNoteData* nd = &(item.second);
+                            if (nd->witnesses.size() > 0) {
+                                nd->witnesses.front().append(note_commitment);
+                            }
+                        }
+                    }
+
+                    // If this is our note, witness it
+                    if (txIsOurs) {
+                        JSOutPoint jsoutpt {hash, i, j};
+                        if (mapWallet[hash].mapNoteData.count(jsoutpt)) {
+                            mapWallet[hash].mapNoteData[jsoutpt].witnesses.push_front(
+                                tree.witness());
+                        }
+                    }
+                }
+            }
+        }
+        vAnchorCache.push_front(tree.root());
+        if (vAnchorCache.size() > WITNESS_CACHE_SIZE) {
+            vAnchorCache.pop_back();
+        }
+        if (fFileBacked) {
+            CWalletDB(strWalletFile).WriteAnchorCache(vAnchorCache);
+        }
+    }
+}
+
+void CWallet::DecrementNoteWitnesses()
+{
+    {
+        LOCK(cs_wallet);
+        for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
+            for (mapNoteData_t::value_type& item : wtxItem.second.mapNoteData) {
+                CNoteData* nd = &(item.second);
+                if (nd->witnesses.size() > 0) {
+                    nd->witnesses.pop_front();
+                }
+            }
+        }
+        if (vAnchorCache.size() > 0) {
+            vAnchorCache.pop_front();
+        }
+        // TODO: If vAnchorCache is empty, we need to regenerate the caches (#1302)
+        if (fFileBacked) {
+            CWalletDB(strWalletFile).WriteAnchorCache(vAnchorCache);
+        }
+    }
+}
+
 bool CWallet::EncryptWallet(const SecureString& strWalletPassphrase)
 {
     if (IsCrypted())
@@ -875,8 +979,9 @@ bool CWallet::AddToWalletIfInvolvingMe(const CTransaction& tx, const CBlock* pbl
         {
             CWalletTx wtx(this,tx);
 
-            if (noteData.size() > 0)
+            if (noteData.size() > 0) {
                 wtx.SetNoteData(noteData);
+            }
 
             // Get merkle branch if transaction was found in a block
             if (pblock)
@@ -965,6 +1070,29 @@ mapNoteData_t CWallet::FindMyNotes(const CTransaction& tx) const
     return noteData;
 }
 
+void CWallet::GetNoteWitnesses(std::vector<JSOutPoint> notes,
+                               std::vector<boost::optional<ZCIncrementalWitness>>& witnesses,
+                               uint256 &final_anchor)
+{
+    {
+        LOCK(cs_wallet);
+        witnesses.resize(notes.size());
+        int i = 0;
+        for (JSOutPoint note : notes) {
+            if (mapWallet.count(note.hash) &&
+                    mapWallet[note.hash].mapNoteData.count(note) &&
+                    mapWallet[note.hash].mapNoteData[note].witnesses.size() > 0) {
+                witnesses[i] = mapWallet[note.hash].mapNoteData[note].witnesses.front();
+            }
+            i++;
+        }
+        // vAnchorCache should only be empty here before the genesis block (so, never)
+        if (vAnchorCache.size() > 0) {
+            final_anchor = vAnchorCache.front();
+        }
+    }
+}
+
 isminetype CWallet::IsMine(const CTxIn &txin) const
 {
     {

src/wallet/wallet.h

@@ -7,6 +7,7 @@
 #define BITCOIN_WALLET_WALLET_H
 
 #include "amount.h"
+#include "coins.h"
 #include "key.h"
 #include "keystore.h"
 #include "primitives/block.h"
@@ -52,6 +53,9 @@ static const unsigned int DEFAULT_TX_CONFIRM_TARGET = 2;
 static const CAmount nHighTransactionMaxFeeWarning = 100 * nHighTransactionFeeWarning;
 //! Largest (in bytes) free transaction we're willing to create
 static const unsigned int MAX_FREE_TRANSACTION_CREATE_SIZE = 1000;
+//! Size of witness cache
+//  Should be large enough that we can expect to never reorg beyond our cache.
+static const unsigned int WITNESS_CACHE_SIZE = 50;
 
 class CAccountingEntry;
 class CBlockIndex;
@@ -201,6 +205,12 @@ public:
     // encrypted.
     uint256 nullifier;
 
+    /**
+     * Cached incremental witnesses for spendable Notes.
+     * Beginning of the list is the most recent witness.
+     */
+    std::list<ZCIncrementalWitness> witnesses;
+
     CNoteData() : address(), nullifier() { }
     CNoteData(libzcash::PaymentAddress a, uint256 n) : address {a}, nullifier {n} { }
 
@@ -568,6 +578,20 @@ private:
     void AddToSpends(const uint256& nullifier, const uint256& wtxid);
     void AddToSpends(const uint256& wtxid);
 
+public:
+    /*
+     * Cached anchors corresponding to the cached incremental witnesses for the
+     * notes in our wallet.
+     */
+    std::list<uint256> vAnchorCache;
+
+protected:
+    void IncrementNoteWitnesses(const CBlockIndex* pindex,
+                                const CBlock* pblock,
+                                const CCoinsViewCache* pcoins);
+    void DecrementNoteWitnesses();
+
+private:
     template <class T>
     void SyncMetaData(std::pair<typename TxSpendMap<T>::iterator, typename TxSpendMap<T>::iterator>);
 
@@ -769,6 +793,10 @@ public:
     std::set<CTxDestination> GetAccountAddresses(std::string strAccount) const;
 
     mapNoteData_t FindMyNotes(const CTransaction& tx) const;
+    void GetNoteWitnesses(
+         std::vector<JSOutPoint> notes,
+         std::vector<boost::optional<ZCIncrementalWitness>>& witnesses,
+         uint256 &final_anchor);
 
     isminetype IsMine(const CTxIn& txin) const;
     CAmount GetDebit(const CTxIn& txin, const isminefilter& filter) const;

src/wallet/walletdb.cpp

@@ -162,6 +162,12 @@ bool CWalletDB::WriteDefaultKey(const CPubKey& vchPubKey)
     return Write(std::string("defaultkey"), vchPubKey);
 }
 
+bool CWalletDB::WriteAnchorCache(const std::list<uint256>& vAnchorCache)
+{
+    nWalletDBUpdated++;
+    return Write(std::string("anchorcache"), vAnchorCache);
+}
+
 bool CWalletDB::ReadPool(int64_t nPool, CKeyPool& keypool)
 {
     return Read(std::make_pair(std::string("pool"), nPool), keypool);
@@ -631,6 +637,10 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
                 return false;
             }
         }
+        else if (strType == "anchorcache")
+        {
+            ssValue >> pwallet->vAnchorCache;
+        }
     } catch (...)
     {
         return false;

src/wallet/walletdb.h

@@ -106,6 +106,8 @@ public:
 
     bool WriteDefaultKey(const CPubKey& vchPubKey);
 
+    bool WriteAnchorCache(const std::list<uint256>& vAnchorCache);
+
     bool ReadPool(int64_t nPool, CKeyPool& keypool);
     bool WritePool(int64_t nPool, const CKeyPool& keypool);
     bool ErasePool(int64_t nPool);