From 63b551378149e135a685df344860cf9426faf2f6 Mon Sep 17 00:00:00 2001
From: fekt <fekt.null@protonmail.com>
Date: Sat, 9 Mar 2024 00:52:29 -0500
Subject: [PATCH] Porting AddToSifted() and AppendSingleSaplingCommitment()

This is still WIP/testing. Not recommended to test with any significant amounts or important wallets.
---
 src/wallet/wallet.cpp | 62 ++++++++++++++++++++++++++++++++++++-------
 src/wallet/wallet.h   |  3 +++
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 5bf0074e7..7dfeb4a70 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -960,6 +960,13 @@ int64_t CWallet::NullifierCount()
     return pcoinsTip->getNullifiers().size();
 }
 
+void CWallet::AddToSifted(const uint256& wtxid)
+{
+    CWalletTx& wtx = mapWallet.at(wtxid);
+
+    if (!wtx.mapSaplingNoteData.empty())
+        setSiftedSapling.emplace(wtxid);
+}
 
 void CWallet::ClearNoteWitnessCache()
 {
@@ -1132,6 +1139,13 @@ int CWallet::VerifyAndSetInitialWitness(const CBlockIndex* pindex, bool witnessO
   return nMinimumHeight;
 }
 
+static SaplingWitness AppendSingleSaplingCommitment(const SaplingWitness& witness, const uint256& commitment)
+{
+    SaplingWitness sw(witness);
+    sw.append(commitment);
+    return sw;
+}
+
 void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
 {
 
@@ -1169,12 +1183,26 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
     pcoinsTip->GetSaplingAnchorAt(saplingRoot, saplingTree);
 
     //Cycle through blocks and transactions building sapling tree until the commitment needed is reached
+    const CBlock *pblock;
     CBlock block;
     if (!ReadBlockFromDisk(block, pblockindex, 1)) {
         throw std::runtime_error(
            strprintf("Cannot read block height %d (%s) from disk", pindex->GetHeight(), pindex->GetBlockHash().GetHex()));
     }
+    else {
+        pblock = &block;
+    }
 
+    std::vector<uint256> vSaplingCommitments;
+    std::vector<SaplingNoteData*> vSaplingNoteData;
+    
+    if (setSiftedSapling.size()) {
+    for (const CTransaction &tx : pblock->vtx) {
+        for (const OutputDescription &outdesc : tx.vShieldedOutput) {
+            vSaplingCommitments.emplace_back(outdesc.cm);
+        }
+    }
+    
     for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
 
       if (wtxItem.second.mapSaplingNoteData.empty())
@@ -1192,18 +1220,29 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
             while (nd->witnesses.size() > WITNESS_CACHE_SIZE) {
                 nd->witnesses.pop_back();
             }
-
-            for (const CTransaction& tx : block.vtx) {
-              for (uint32_t i = 0; i < tx.vShieldedOutput.size(); i++) {
-                const uint256& note_commitment = tx.vShieldedOutput[i].cm;
-                nd->witnesses.front().append(note_commitment);
-              }
-            }
+            
+            // for async append of commitments
+            vSaplingNoteData.push_back(nd);
+            
             nd->witnessHeight = pblockindex->GetHeight();
           }
         }
 
       }
+        // Parallelization with std::async
+        std::vector<std::future<SaplingWitness>> vSaplingWitnessFutures;
+        for (const auto& commitment : vSaplingCommitments) {
+            for (auto pnd : vSaplingNoteData) {
+                vSaplingWitnessFutures.emplace_back(std::async(std::launch::async, AppendSingleSaplingCommitment, pnd->witnesses.front(), commitment));
+            }
+            assert(vSaplingWitnessFutures.size() == vSaplingNoteData.size());
+
+            for (int i = 0; i < vSaplingWitnessFutures.size(); i++) {
+                vSaplingNoteData.at(i)->witnesses.front() = vSaplingWitnessFutures.at(i).get();
+            }
+
+            vSaplingWitnessFutures.resize(0);
+        }
     }
 
     if (pblockindex == pindex)
@@ -1212,7 +1251,7 @@ void CWallet::BuildWitnessCache(const CBlockIndex* pindex, bool witnessOnly)
     pblockindex = chainActive.Next(pblockindex);
 
   }
-
+ }
 }
 
 bool CWallet::EncryptWallet(const SecureString& strWalletPassphrase)
@@ -1466,6 +1505,7 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletD
         mapWallet[hash].BindWallet(this);
         UpdateNullifierNoteMapWithTx(mapWallet[hash]);
         AddToSpends(hash);
+        AddToSifted(hash);
     }
     else
     {
@@ -1558,9 +1598,13 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletD
 
         // Write to disk
         if (fInsertedNew || fUpdated)
+        {
+            AddToSifted(hash);
+         
             if (!wtx.WriteToDisk(pwalletdb))
                 return false;
-
+        }
+        
         // Break debit/credit balance caches:
         wtx.MarkDirty();
 
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 39b52c62c..606aa243d 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -1015,6 +1015,9 @@ public:
     std::map<uint256, SaplingOutPoint> mapSaplingNullifiersToNotes;
     std::map<uint256, CWalletTx> mapWallet;
 
+    std::set<uint256> setSiftedSapling;
+    void AddToSifted(const uint256& wtxid);
+    
     int64_t nOrderPosNext;
 
     std::map<CTxDestination, CAddressBookData> mapAddressBook;