Attempt to port UpdateSpentHeightAndMaybePruneWitnesses

src/wallet/wallet.cpp

@@ -1103,7 +1103,7 @@ int CWallet::VerifyAndSetInitialWitness(const CBlockIndex* pindex, bool witnessO
               nd->witnesses.front().append(note_commitment);
             }
 
-            //Only needed for intial witness
+            //Only needed for initial witness
             if (nd->witnesses.empty()) {
               saplingTree.append(note_commitment);
 
@@ -1140,6 +1140,7 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
     return;
   }
 
+  std::vector<uint256> nullifiers;
   uint256 saplingRoot;
   CBlockIndex* pblockindex = chainActive[startHeight];
   int height = chainActive.Height();
@@ -1172,16 +1173,42 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
            strprintf("Cannot read block height %d (%s) from disk", pindex->GetHeight(), pindex->GetBlockHash().GetHex()));
     }
 
+    // block height of the current block being processed
+    int indexHeight = pblockindex->GetHeight();
     for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
 
       if (wtxItem.second.mapSaplingNoteData.empty())
         continue;
 
+      // ignore any tx that is not yet in a block (i.e. orphaned or in the mempool)
       if (wtxItem.second.GetDepthInMainChain() > 0) {
-
-        //Sapling
         for (mapSaplingNoteData_t::value_type& item : wtxItem.second.mapSaplingNoteData) {
           auto* nd = &(item.second);
+
+          // update spentness info
+          for (const auto& nullifier : nullifiers) {
+            // This is UpdateSpentHeightAndMaybePruneWitnesses ported to our codebase
+            // If the note has no witnesses, then either the note has not been mined
+            // (and thus cannot be spent at this height), or has been spent for long
+            // enough that we will never unspend it. Either way, we can skip the
+            // spentness check and pruning.
+            if (nd->witnesses.empty()) continue;
+            // Update spentness for sapling note data. If a note is in a block and has a nullifier
+            // it is a spend
+            if (nd->nullifier.has_value() && nd->nullifier.value() == nullifier) {
+                LogPrintf("%s: updating spentness for %s spentHeight=%d\n", __func__, nullifier.GetHex().c_str(), nd->spentHeight);
+                nd->spentHeight = indexHeight;
+            }
+            // Prune witnesses for notes spent more than WITNESS_CACHE_SIZE blocks ago,
+            // so we stop updating their witnesses. This is safe to do because we know
+            // we won't roll back more than WITNESS_CACHE_SIZE blocks, since it is a
+            // number larger than the max reorg length
+            if (nd->spentHeight.has_value() && nd->spentHeight.value() + WITNESS_CACHE_SIZE < indexHeight) {
+                LogPrintf("%s: pruning witnesses for %s at indexHeight=%d with spentHeight=%d\n", __func__, nullifier.GetHex().c_str(), indexHeight, nd->spentHeight);
+                nd->witnesses.clear();
+                nd->witnessHeight = -1;
+            }
+          }
           if (nd->nullifier && nd->witnessHeight == pblockindex->GetHeight() - 1
               && GetSaplingSpendDepth(*item.second.nullifier) <= WITNESS_CACHE_SIZE) {
 
@@ -1195,6 +1222,11 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
                 const uint256& note_commitment = tx.vShieldedOutput[i].cm;
                 nd->witnesses.front().append(note_commitment);
               }
+              // add all nullifiers in this tx to our list
+              for (uint32_t i = 0; i < tx.vShieldedSpend.size(); i++) {
+                //TODO: do we need to make sure it's not already there before adding it?
+                nullifiers.emplace_back(tx.vShieldedSpend[i].nullifier);
+              }
             }
             nd->witnessHeight = pblockindex->GetHeight();
           }
@@ -1207,9 +1239,7 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
       break;
 
     pblockindex = chainActive.Next(pblockindex);
-
   }
-
 }
 
 bool CWallet::EncryptWallet(const SecureString& strWalletPassphrase)