Correctly set CNoteData::witnessHeight when decrementing witness caches

Closes #1715

src/wallet/gtest/test_wallet.cpp

@@ -50,8 +50,8 @@ public:
                                 ZCIncrementalMerkleTree tree) {
         CWallet::IncrementNoteWitnesses(pindex, pblock, tree);
     }
-    void DecrementNoteWitnesses() {
+    void DecrementNoteWitnesses(const CBlockIndex* pindex) {
-        CWallet::DecrementNoteWitnesses();
+        CWallet::DecrementNoteWitnesses(pindex);
     }
     void WriteWitnessCache(MockWalletDB& walletdb) {
         CWallet::WriteWitnessCache(walletdb);
@@ -675,7 +675,7 @@ TEST(wallet_tests, cached_witnesses_empty_chain) {
     EXPECT_TRUE((bool) witnesses[1]);
 
     // Until #1302 is implemented, this should triggger an assertion
-    EXPECT_DEATH(wallet.DecrementNoteWitnesses(),
+    EXPECT_DEATH(wallet.DecrementNoteWitnesses(&index),
                  "Assertion `nWitnessCacheSize > 0' failed.");
 }
 
@@ -748,7 +748,7 @@ TEST(wallet_tests, cached_witnesses_chain_tip) {
 
         // Decrementing should give us the previous anchor
         uint256 anchor3;
-        wallet.DecrementNoteWitnesses();
+        wallet.DecrementNoteWitnesses(&index2);
         witnesses.clear();
         wallet.GetNoteWitnesses(notes, witnesses, anchor3);
         EXPECT_FALSE((bool) witnesses[0]);
@@ -773,6 +773,101 @@ TEST(wallet_tests, cached_witnesses_chain_tip) {
     }
 }
 
+TEST(wallet_tests, CachedWitnessesDecrementFirst) {
+    TestWallet wallet;
+    uint256 anchor2;
+    CBlock block2;
+    CBlockIndex index2(block2);
+    ZCIncrementalMerkleTree tree;
+
+    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.GetHash(), 0, 1};
+        CNoteData nd {sk.address(), nullifier};
+        noteData[jsoutpt] = nd;
+        wtx.SetNoteData(noteData);
+        wallet.AddToWallet(wtx, true, NULL);
+
+        // First block (case tested in _empty_chain)
+        CBlock block1;
+        block1.vtx.push_back(wtx);
+        CBlockIndex index1(block1);
+        index1.nHeight = 1;
+        wallet.IncrementNoteWitnesses(&index1, &block1, tree);
+    }
+
+    {
+        // Second transaction (case tested in _chain_tip)
+        auto wtx = GetValidReceive(sk, 50, true);
+        auto note = GetNote(sk, wtx, 0, 1);
+        auto nullifier = note.nullifier(sk);
+
+        mapNoteData_t noteData;
+        JSOutPoint jsoutpt {wtx.GetHash(), 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;
+
+        // Second block (case tested in _chain_tip)
+        block2.vtx.push_back(wtx);
+        index2.nHeight = 2;
+        wallet.IncrementNoteWitnesses(&index2, &block2, tree);
+        // Called to fetch anchor
+        wallet.GetNoteWitnesses(notes, witnesses, anchor2);
+    }
+
+    {
+        // Third transaction - never mined
+        auto wtx = GetValidReceive(sk, 20, true);
+        auto note = GetNote(sk, wtx, 0, 1);
+        auto nullifier = note.nullifier(sk);
+
+        mapNoteData_t noteData;
+        JSOutPoint jsoutpt {wtx.GetHash(), 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 anchor3;
+
+        wallet.GetNoteWitnesses(notes, witnesses, anchor3);
+        EXPECT_FALSE((bool) witnesses[0]);
+
+        // Decrementing (before the transaction has ever seen an increment)
+        // should give us the previous anchor
+        uint256 anchor4;
+        wallet.DecrementNoteWitnesses(&index2);
+        witnesses.clear();
+        wallet.GetNoteWitnesses(notes, witnesses, anchor4);
+        EXPECT_FALSE((bool) witnesses[0]);
+        // Should not equal second anchor because none of these notes had witnesses
+        EXPECT_NE(anchor2, anchor4);
+
+        // Re-incrementing with the same block should give the same result
+        uint256 anchor5;
+        wallet.IncrementNoteWitnesses(&index2, &block2, tree);
+        witnesses.clear();
+        wallet.GetNoteWitnesses(notes, witnesses, anchor5);
+        EXPECT_FALSE((bool) witnesses[0]);
+        EXPECT_EQ(anchor3, anchor5);
+    }
+}
+
 TEST(wallet_tests, ClearNoteWitnessCache) {
     TestWallet wallet;
 

src/wallet/wallet.cpp

@@ -373,7 +373,7 @@ void CWallet::ChainTip(const CBlockIndex *pindex, const CBlock *pblock,
     if (added) {
         IncrementNoteWitnesses(pindex, pblock, tree);
     } else {
-        DecrementNoteWitnesses();
+        DecrementNoteWitnesses(pindex);
     }
 }
 
@@ -648,7 +648,7 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                 // Only increment witnesses that are behind the current height
                 if (nd->witnessHeight < pindex->nHeight) {
                     // Witnesses being incremented should always be either -1
-                    // (never incremented) or one below pindex
+                    // (never incremented or decremented) or one below pindex
                     assert((nd->witnessHeight == -1) ||
                            (nd->witnessHeight == pindex->nHeight - 1));
                     // Copy the witness for the previous block if we have one
@@ -746,7 +746,7 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
     }
 }
 
-void CWallet::DecrementNoteWitnesses()
+void CWallet::DecrementNoteWitnesses(const CBlockIndex* pindex)
 {
     {
         LOCK(cs_wallet);
@@ -755,10 +755,14 @@ void CWallet::DecrementNoteWitnesses()
                 CNoteData* nd = &(item.second);
                 // Check the validity of the cache
                 assert(nWitnessCacheSize >= nd->witnesses.size());
+                // Witnesses being decremented should always be either -1
+                // (never incremented or decremented) or equal to pindex
+                assert((nd->witnessHeight == -1) ||
+                       (nd->witnessHeight == pindex->nHeight));
                 if (nd->witnesses.size() > 0) {
                     nd->witnesses.pop_front();
                 }
-                nd->witnessHeight -= 1;
+                nd->witnessHeight = pindex->nHeight - 1;
             }
         }
         nWitnessCacheSize -= 1;

src/wallet/wallet.h

@@ -619,7 +619,7 @@ protected:
     void IncrementNoteWitnesses(const CBlockIndex* pindex,
                                 const CBlock* pblock,
                                 ZCIncrementalMerkleTree& tree);
-    void DecrementNoteWitnesses();
+    void DecrementNoteWitnesses(const CBlockIndex* pindex);
 
     template <typename WalletDB>
     void WriteWitnessCache(WalletDB& walletdb) {