Fixing bug at batch_container which resulted fails in case we had reo…

…rgs in chain
firoorg · Jun 30, 2021 · f4a4cba · f4a4cba
1 parent 27bc584
commit f4a4cba
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 58 deletions.
diff --git a/src/batchproof_container.cpp b/src/batchproof_container.cpp
@@ -69,57 +69,50 @@ void BatchProofContainer::add(lelantus::JoinSplit* joinSplit,
 
 void BatchProofContainer::removeSigma(const sigma::spend_info_container& spendSerials) {
     for (auto& spendSerial : spendSerials) {
-        bool foundAtSigma =  false;
         for (auto& itr :sigmaProofs) {
             if (itr.first.first == spendSerial.second.denomination && itr.first.second.first == spendSerial.second.coinGroupId) {
                 auto& vProofs = itr.second;
                 for (auto dataItr = vProofs.begin(); dataItr != vProofs.end(); dataItr++) {
                     if (dataItr->coinSerialNumber == spendSerial.first) {
                         vProofs.erase(dataItr);
-                        foundAtSigma = true;
                     }
                 }
             }
         }
-        if (!foundAtSigma) {
-            int64_t denom;
-            sigma::DenominationToInteger(spendSerial.second.denomination, denom);
-            int id = denom / 1000 + spendSerial.second.coinGroupId;
-            // afterFixes bool with the pair of set id is considered separate set identifiers, so try to find in one set, if not found try also in another
-            std::pair<std::pair<uint32_t, bool>, bool> key1 = std::make_pair(std::make_pair(id, false), true);
-            std::pair<std::pair<uint32_t, bool>, bool> key2 = std::make_pair(std::make_pair(id, true), true);
-            std::vector<LelantusSigmaProofData>* vProofs;
-            if (lelantusSigmaProofs.count(key1) > 0)
-                vProofs = &lelantusSigmaProofs[key1];
-            else if (lelantusSigmaProofs.count(key2) > 0)
-                vProofs = &lelantusSigmaProofs[key2];
-            else
-                continue;
-            for (auto dataItr = vProofs->begin(); dataItr != vProofs->end(); dataItr++) {
-                if (dataItr->serialNumber == spendSerial.first) {
-                    vProofs->erase(dataItr);
-                }
-            }
-        }
     }
 }
 void BatchProofContainer::removeLelantus(std::unordered_map<Scalar, int> spentSerials) {
     for (auto& spendSerial : spentSerials) {
+
+        int id = spendSerial.second;
+        int coinGroupId = id % (CENT / 1000);
+        int64_t intDenom = (id - coinGroupId);
+        intDenom *= 1000;
+        sigma::CoinDenomination denomination;
+        bool isSigmaToLela = false;
+        if (sigma::IntegerToDenomination(intDenom, denomination))
+            isSigmaToLela = true;
+
         // afterFixes bool with the pair of set id is considered separate set identifiers, so try to find in one set, if not found try also in another
-        std::pair<std::pair<uint32_t, bool>, bool> key1 = std::make_pair(std::make_pair(spendSerial.second, false), true);
-        std::pair<std::pair<uint32_t, bool>, bool> key2 = std::make_pair(std::make_pair(spendSerial.second, true), true);
+        std::pair<std::pair<uint32_t, bool>, bool> key1 = std::make_pair(std::make_pair(id, false), isSigmaToLela);
+        std::pair<std::pair<uint32_t, bool>, bool> key2 = std::make_pair(std::make_pair(id, true), isSigmaToLela);
         std::vector<LelantusSigmaProofData>* vProofs;
-        if (lelantusSigmaProofs.count(key1) > 0)
+        if (lelantusSigmaProofs.count(key1) > 0) {
             vProofs = &lelantusSigmaProofs[key1];
-        else if (lelantusSigmaProofs.count(key2) > 0)
+            erase(vProofs, spendSerial.first);
+        }
+
+        if (lelantusSigmaProofs.count(key2) > 0) {
             vProofs = &lelantusSigmaProofs[key2];
-        else
-            continue;
+            erase(vProofs, spendSerial.first);
+        }
+    }
+}
 
-        for (auto dataItr = vProofs->begin(); dataItr != vProofs->end(); dataItr++) {
-            if (dataItr->serialNumber == spendSerial.first) {
-                vProofs->erase(dataItr);
-            }
+void BatchProofContainer::erase(std::vector<LelantusSigmaProofData>* vProofs, const Scalar& serial) {
+    for (auto dataItr = vProofs->begin(); dataItr != vProofs->end(); dataItr++) {
+        if (dataItr->serialNumber == serial) {
+            vProofs->erase(dataItr);
         }
     }
 }

diff --git a/src/batchproof_container.h b/src/batchproof_container.h
@@ -29,6 +29,7 @@ class BatchProofContainer {
 
     void removeSigma(const sigma::spend_info_container& spendSerials);
     void removeLelantus(std::unordered_map<Scalar, int> spentSerials);
+    void erase(std::vector<LelantusSigmaProofData>* vProofs, const Scalar& serial);
 
     void batch_sigma();
     void batch_lelantus();

diff --git a/src/validation.cpp b/src/validation.cpp
@@ -3194,7 +3194,6 @@ bool static DisconnectTip(CValidationState& state, const CChainParams& chainpara
     for (CTransactionRef tx : block.vtx) {
         CheckTransaction(*tx, state, false, tx->GetHash(), false, pindexDelete->pprev->nHeight,
             false, false, block.sigmaTxInfo.get(), block.lelantusTxInfo.get());
-
         if(GetBoolArg("-batching", true)) {
             if (tx->IsLelantusJoinSplit()) {
                 const CTxIn &txin = tx->vin[0];
@@ -3214,22 +3213,8 @@ bool static DisconnectTip(CValidationState& state, const CChainParams& chainpara
                     continue;
                 }
 
-                if (joinsplit->isSigmaToLelantus()) {
-                    for (size_t i = 0; i < serials.size(); i++) {
-                        int coinGroupId = ids[i] % (CENT / 1000);
-                        int64_t intDenom = (ids[i] - coinGroupId);
-                        intDenom *= 1000;
-                        sigma::CoinDenomination denomination;
-                        if (!sigma::IntegerToDenomination(intDenom, denomination))
-                            lelantusSerialsToRemove.insert(std::make_pair(serials[i], ids[i]));
-                        else
-                            sigmaSerialsToRemove.insert(std::make_pair(
-                                    serials[i], sigma::CSpendCoinInfo::make(denomination, coinGroupId)));
-                    }
-                } else {
-                    for (size_t i = 0; i < serials.size(); i++) {
-                        lelantusSerialsToRemove.insert(std::make_pair(serials[i], ids[i]));
-                    }
+                for (size_t i = 0; i < serials.size(); i++) {
+                    lelantusSerialsToRemove.insert(std::make_pair(serials[i], ids[i]));
                 }
             } else if (tx->IsSigmaSpend()) {
                 for (const CTxIn &txin : tx->vin) {
@@ -3267,6 +3252,16 @@ bool static DisconnectTip(CValidationState& state, const CChainParams& chainpara
 
 	sigma::DisconnectTipSigma(block, pindexDelete);
     lelantus::DisconnectTipLelantus(block, pindexDelete);
+
+    BatchProofContainer* batchProofContainer = BatchProofContainer::get_instance();
+    if (sigmaSerialsToRemove.size() > 0) {
+        batchProofContainer->removeSigma(sigmaSerialsToRemove);
+    }
+
+    if (lelantusSerialsToRemove.size() > 0) {
+        batchProofContainer->removeLelantus(lelantusSerialsToRemove);
+    }
+
     // Roll back MTP state
     MTPState::GetMTPState()->SetLastBlock(pindexDelete->pprev, chainparams.GetConsensus());
 
@@ -3313,15 +3308,6 @@ bool static DisconnectTip(CValidationState& state, const CChainParams& chainpara
     // Update chainActive and related variables.
     UpdateTip(pindexDelete->pprev, chainparams);
 
-    BatchProofContainer* batchProofContainer = BatchProofContainer::get_instance();
-    if (sigmaSerialsToRemove.size() > 0) {
-        batchProofContainer->removeSigma(sigmaSerialsToRemove);
-    }
-
-    if (lelantusSerialsToRemove.size() > 0) {
-        batchProofContainer->removeLelantus(lelantusSerialsToRemove);
-    }
-
 #ifdef ENABLE_WALLET
     // update mint/spend wallet
     if (!GetBoolArg("-disablewallet", false) && pwalletMain->zwallet) {