Increment witnesses for new transactions on rescan

Closes #1400

src/chain.h

@@ -138,6 +138,9 @@ public:
     //! Verification status of this block. See enum BlockStatus
     unsigned int nStatus;
 
+    //! The anchor for the tree state up to the start of this block
+    uint256 hashAnchor;
+
     //! block header
     int nVersion;
     uint256 hashMerkleRoot;
@@ -163,6 +166,7 @@ public:
         nTx = 0;
         nChainTx = 0;
         nStatus = 0;
+        hashAnchor = uint256();
         nSequenceId = 0;
 
         nVersion       = 0;
@@ -320,6 +324,7 @@ public:
             READWRITE(VARINT(nDataPos));
         if (nStatus & BLOCK_HAVE_UNDO)
             READWRITE(VARINT(nUndoPos));
+        READWRITE(hashAnchor);
 
         // block header
         READWRITE(this->nVersion);

src/main.cpp

@@ -2079,6 +2079,8 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
     // Construct the incremental merkle tree at the current
     // block position,
     auto old_tree_root = view.GetBestAnchor();
+    // saving the top anchor in the block index as we go.
+    pindex->hashAnchor = old_tree_root;
     ZCIncrementalMerkleTree tree;
     // This should never fail: we should always be able to get the root
     // that is on the tip of our chain

src/txdb.cpp

@@ -295,6 +295,7 @@ bool CBlockTreeDB::LoadBlockIndexGuts()
                 pindexNew->nFile          = diskindex.nFile;
                 pindexNew->nDataPos       = diskindex.nDataPos;
                 pindexNew->nUndoPos       = diskindex.nUndoPos;
+                pindexNew->hashAnchor     = diskindex.hashAnchor;
                 pindexNew->nVersion       = diskindex.nVersion;
                 pindexNew->hashMerkleRoot = diskindex.hashMerkleRoot;
                 pindexNew->nTime          = diskindex.nTime;

src/wallet/gtest/test_wallet.cpp

@@ -638,8 +638,9 @@ TEST(wallet_tests, cached_witnesses_empty_chain) {
 
     CBlock block;
     block.vtx.push_back(wtx);
+    CBlockIndex index(block);
     ZCIncrementalMerkleTree tree;
-    wallet.IncrementNoteWitnesses(NULL, &block, tree);
+    wallet.IncrementNoteWitnesses(&index, &block, tree);
     witnesses.clear();
     wallet.GetNoteWitnesses(notes, witnesses, anchor);
     EXPECT_TRUE((bool) witnesses[0]);
@@ -677,7 +678,9 @@ TEST(wallet_tests, cached_witnesses_chain_tip) {
 
         // First block (case tested in _empty_chain)
         block1.vtx.push_back(wtx);
-        wallet.IncrementNoteWitnesses(NULL, &block1, tree);
+        CBlockIndex index1(block1);
+        index1.nHeight = 1;
+        wallet.IncrementNoteWitnesses(&index1, &block1, tree);
         // Called to fetch anchor
         wallet.GetNoteWitnesses(notes, witnesses, anchor1);
     }
@@ -706,8 +709,10 @@ TEST(wallet_tests, cached_witnesses_chain_tip) {
         CBlock block2;
         block2.hashPrevBlock = block1.GetHash();
         block2.vtx.push_back(wtx);
+        CBlockIndex index2(block2);
+        index2.nHeight = 2;
         ZCIncrementalMerkleTree tree2 {tree};
-        wallet.IncrementNoteWitnesses(NULL, &block2, tree2);
+        wallet.IncrementNoteWitnesses(&index2, &block2, tree2);
         witnesses.clear();
         wallet.GetNoteWitnesses(notes, witnesses, anchor2);
         EXPECT_TRUE((bool) witnesses[0]);
@@ -724,7 +729,7 @@ TEST(wallet_tests, cached_witnesses_chain_tip) {
 
         // Re-incrementing with the same block should give the same result
         uint256 anchor4;
-        wallet.IncrementNoteWitnesses(NULL, &block2, tree);
+        wallet.IncrementNoteWitnesses(&index2, &block2, tree);
         witnesses.clear();
         wallet.GetNoteWitnesses(notes, witnesses, anchor4);
         EXPECT_TRUE((bool) witnesses[0]);

src/wallet/wallet.cpp

@@ -636,7 +636,7 @@ void CWallet::ClearNoteWitnessCache()
 
 void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                                      const CBlock* pblockIn,
-                                     ZCIncrementalMerkleTree tree)
+                                     ZCIncrementalMerkleTree& tree)
 {
     {
         LOCK(cs_wallet);
@@ -645,12 +645,19 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                 CNoteData* nd = &(item.second);
                 // Check the validity of the cache
                 assert(nWitnessCacheSize >= nd->witnesses.size());
-                // Copy the witness for the previous block if we have one
+                // Only increment witnesses that are behind the current height
-                if (nd->witnesses.size() > 0) {
+                if (nd->witnessHeight < pindex->nHeight) {
-                    nd->witnesses.push_front(nd->witnesses.front());
+                    // Witnesses being incremented should always be either -1
-                }
+                    // (never incremented) or one below pindex
-                if (nd->witnesses.size() > WITNESS_CACHE_SIZE) {
+                    assert((nd->witnessHeight == -1) ||
-                    nd->witnesses.pop_back();
+                           (nd->witnessHeight == pindex->nHeight - 1));
+                    // 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();
+                    }
                 }
             }
         }
@@ -680,7 +687,8 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                             CNoteData* nd = &(item.second);
                             // Check the validity of the cache
                             assert(nWitnessCacheSize >= nd->witnesses.size());
-                            if (nd->witnesses.size() > 0) {
+                            if (nd->witnessHeight < pindex->nHeight &&
+                                    nd->witnesses.size() > 0) {
                                 nd->witnesses.front().append(note_commitment);
                             }
                         }
@@ -693,6 +701,8 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                             CNoteData* nd = &(mapWallet[hash].mapNoteData[jsoutpt]);
                             assert(nd->witnesses.size() == 0);
                             nd->witnesses.push_front(tree.witness());
+                            // Set height to one less than pindex so it gets incremented
+                            nd->witnessHeight = pindex->nHeight - 1;
                             // Check the validity of the cache
                             assert(nWitnessCacheSize >= nd->witnesses.size());
                         }
@@ -700,13 +710,19 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                 }
             }
         }
+
+        // Update witness heights
         for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
             for (mapNoteData_t::value_type& item : wtxItem.second.mapNoteData) {
                 CNoteData* nd = &(item.second);
+                if (nd->witnessHeight < pindex->nHeight) {
+                    nd->witnessHeight = pindex->nHeight;
+                }
                 // Check the validity of the cache
                 assert(nWitnessCacheSize >= nd->witnesses.size());
             }
         }
+
         if (fFileBacked) {
             CWalletDB walletdb(strWalletFile);
             WriteWitnessCache(walletdb);
@@ -726,6 +742,7 @@ void CWallet::DecrementNoteWitnesses()
                 if (nd->witnesses.size() > 0) {
                     nd->witnesses.pop_front();
                 }
+                nd->witnessHeight -= 1;
             }
         }
         nWitnessCacheSize -= 1;
@@ -1716,6 +1733,14 @@ int CWallet::ScanForWalletTransactions(CBlockIndex* pindexStart, bool fUpdate)
                 if (AddToWalletIfInvolvingMe(tx, &block, fUpdate))
                     ret++;
             }
+
+            ZCIncrementalMerkleTree tree;
+            // This should never fail: we should always be able to get the tree
+            // state on the path to the tip of our chain
+            assert(pcoinsTip->GetAnchorAt(pindex->hashAnchor, tree));
+            // Increment note witness caches
+            IncrementNoteWitnesses(pindex, &block, tree);
+
             pindex = chainActive.Next(pindex);
             if (GetTime() >= nNow + 60) {
                 nNow = GetTime();

src/wallet/wallet.h

@@ -221,9 +221,14 @@ public:
      */
     std::list<ZCIncrementalWitness> witnesses;
 
-    CNoteData() : address(), nullifier() { }
+    /** Block height corresponding to the most current witness. */
-    CNoteData(libzcash::PaymentAddress a) : address {a}, nullifier() { }
+    int witnessHeight;
-    CNoteData(libzcash::PaymentAddress a, uint256 n) : address {a}, nullifier {n} { }
+
+    CNoteData() : address(), nullifier(), witnessHeight {-1} { }
+    CNoteData(libzcash::PaymentAddress a) :
+            address {a}, nullifier(), witnessHeight {-1} { }
+    CNoteData(libzcash::PaymentAddress a, uint256 n) :
+            address {a}, nullifier {n}, witnessHeight {-1} { }
 
     ADD_SERIALIZE_METHODS;
 
@@ -232,6 +237,7 @@ public:
         READWRITE(address);
         READWRITE(nullifier);
         READWRITE(witnesses);
+        READWRITE(witnessHeight);
     }
 
     friend bool operator<(const CNoteData& a, const CNoteData& b) {
@@ -612,7 +618,7 @@ public:
 protected:
     void IncrementNoteWitnesses(const CBlockIndex* pindex,
                                 const CBlock* pblock,
-                                ZCIncrementalMerkleTree tree);
+                                ZCIncrementalMerkleTree& tree);
     void DecrementNoteWitnesses();
 
     template <typename WalletDB>