diff --git a/src/avalanche/proofpool.cpp b/src/avalanche/proofpool.cpp index 288dc3dc7..8ff71fa44 100644 --- a/src/avalanche/proofpool.cpp +++ b/src/avalanche/proofpool.cpp @@ -1,56 +1,62 @@ // Copyright (c) 2021 The Bitcoin developers // Distributed under the MIT software license, see the accompanying // file COPYING or http://www.opensource.org/licenses/mit-license.php. #include namespace avalanche { ProofPool::AddProofStatus ProofPool::addProof(const ProofRef &proof) { assert(proof); const ProofId &proofid = proof->getId(); auto &poolView = pool.get(); if (poolView.find(proofid) != poolView.end()) { return AddProofStatus::DUPLICATED; } // Attach UTXOs to this proof. std::unordered_set conflicting_proofs; for (size_t i = 0; i < proof->getStakes().size(); i++) { auto p = pool.emplace(i, proof); if (!p.second) { // We have a collision with an existing proof. conflicting_proofs.insert(p.first->proof); } } // For now, if there is a conflict, just cleanup the mess. if (conflicting_proofs.size() > 0) { for (const auto &s : proof->getStakes()) { auto it = pool.find(s.getStake().getUTXO()); assert(it != pool.end()); // We need to delete that one. if (it->proof->getId() == proofid) { pool.erase(it); } } return AddProofStatus::REJECTED; } return AddProofStatus::SUCCEED; } // Having the ProofRef passed by reference is risky because the proof could be // deleted during the erasure loop, so we pass it by value. Since it's a shared // pointer, the copy is cheap enough and should not have any significant impact // on performance. bool ProofPool::removeProof(ProofRef proof) { auto &poolView = pool.get(); return poolView.erase(proof->getId()); } +ProofRef ProofPool::getProof(const ProofId &proofid) const { + auto &poolView = pool.get(); + auto it = poolView.find(proofid); + return it == poolView.end() ? nullptr : it->proof; +} + } // namespace avalanche diff --git a/src/avalanche/proofpool.h b/src/avalanche/proofpool.h index c7de6e0f8..0f206dac3 100644 --- a/src/avalanche/proofpool.h +++ b/src/avalanche/proofpool.h @@ -1,76 +1,78 @@ // Copyright (c) 2021 The Bitcoin developers // Distributed under the MIT software license, see the accompanying // file COPYING or http://www.opensource.org/licenses/mit-license.php. #ifndef BITCOIN_AVALANCHE_PROOFPOOL_H #define BITCOIN_AVALANCHE_PROOFPOOL_H #include #include #include #include #include #include #include #include namespace avalanche { struct ProofPoolEntry { size_t utxoIndex; ProofRef proof; const COutPoint &getUTXO() const { return proof->getStakes().at(utxoIndex).getStake().getUTXO(); } ProofPoolEntry(size_t _utxoIndex, ProofRef _proof) : utxoIndex(_utxoIndex), proof(std::move(_proof)) {} }; struct by_utxo; struct by_proofid; struct ProofPoolEntryProofIdKeyExtractor { using result_type = ProofId; result_type operator()(const ProofPoolEntry &entry) const { return entry.proof->getId(); } }; namespace bmi = boost::multi_index; /** * Map a proof to each utxo. A proof can be mapped with several utxos. */ struct ProofPool { boost::multi_index_container< ProofPoolEntry, bmi::indexed_by< // index by utxo bmi::hashed_unique< bmi::tag, bmi::const_mem_fun, SaltedOutpointHasher>, // index by proofid bmi::hashed_non_unique, ProofPoolEntryProofIdKeyExtractor, SaltedProofIdHasher>>> pool; enum AddProofStatus { REJECTED = 0, //!< Rejected due to conflicts SUCCEED = 1, //!< Added successfully DUPLICATED = 2, //!< Already in pool }; AddProofStatus addProof(const ProofRef &proof); bool removeProof(ProofRef proof); + + ProofRef getProof(const ProofId &proofid) const; }; } // namespace avalanche #endif // BITCOIN_AVALANCHE_PROOFPOOL_H diff --git a/src/avalanche/test/proofpool_tests.cpp b/src/avalanche/test/proofpool_tests.cpp index 7cf060aee..172783186 100644 --- a/src/avalanche/test/proofpool_tests.cpp +++ b/src/avalanche/test/proofpool_tests.cpp @@ -1,70 +1,88 @@ // Copyright (c) 2021 The Bitcoin developers // Distributed under the MIT software license, see the accompanying // file COPYING or http://www.opensource.org/licenses/mit-license.php. #include #include #include #include #include #include #include #include using namespace avalanche; BOOST_FIXTURE_TEST_SUITE(proofpool_tests, TestingSetup) BOOST_AUTO_TEST_CASE(add_remove_proof) { ProofPool testPool; std::vector proofs; for (size_t i = 0; i < 10; i++) { // Add a bunch of random proofs auto proof = buildRandomProof(MIN_VALID_PROOF_SCORE); BOOST_CHECK_EQUAL(testPool.addProof(proof), ProofPool::AddProofStatus::SUCCEED); // Trying to add them again will return a duplicated status for (size_t j = 0; j < 10; j++) { BOOST_CHECK_EQUAL(testPool.addProof(proof), ProofPool::AddProofStatus::DUPLICATED); } proofs.push_back(std::move(proof)); } const CKey key = CKey::MakeCompressedKey(); const COutPoint conflictingOutpoint{TxId(GetRandHash()), 0}; auto buildProofWithSequence = [&](uint64_t sequence) { ProofBuilder pb(sequence, 0, key); BOOST_CHECK( pb.addUTXO(conflictingOutpoint, 10 * COIN, 123456, false, key)); return pb.build(); }; auto proof_seq10 = buildProofWithSequence(10); BOOST_CHECK_EQUAL(testPool.addProof(proof_seq10), ProofPool::AddProofStatus::SUCCEED); proofs.push_back(std::move(proof_seq10)); auto proof_seq20 = buildProofWithSequence(20); BOOST_CHECK_EQUAL(testPool.addProof(proof_seq20), ProofPool::AddProofStatus::REJECTED); // Removing proofs which are not in the pool will fail for (size_t i = 0; i < 10; i++) { BOOST_CHECK( !testPool.removeProof(buildRandomProof(MIN_VALID_PROOF_SCORE))); } for (auto proof : proofs) { BOOST_CHECK(testPool.removeProof(proof)); } BOOST_CHECK(testPool.pool.empty()); } +BOOST_AUTO_TEST_CASE(get_proof) { + ProofPool testPool; + + for (size_t i = 0; i < 10; i++) { + BOOST_CHECK(!testPool.getProof(ProofId(GetRandHash()))); + } + + for (size_t i = 0; i < 10; i++) { + auto proof = buildRandomProof(MIN_VALID_PROOF_SCORE); + BOOST_CHECK_EQUAL(testPool.addProof(proof), + ProofPool::AddProofStatus::SUCCEED); + + auto retrievedProof = testPool.getProof(proof->getId()); + BOOST_CHECK_NE(retrievedProof, nullptr); + BOOST_CHECK_EQUAL(retrievedProof->getId(), proof->getId()); + } +} + BOOST_AUTO_TEST_SUITE_END()