From b7b2288a746f7510d3cba667dc53728407acc804 Mon Sep 17 00:00:00 2001
From: Byron Hambly <bhambly@blockstream.com>
Date: Thu, 9 May 2024 10:02:57 +0200
Subject: [PATCH] discount: implement mempool logic

---
 src/core_write.cpp  |  5 ++++
 src/policy/policy.h |  1 +
 src/txmempool.cpp   | 37 +++++++++++++++++++++-------
 src/txmempool.h     | 60 ++++++++++++++++++++++++++++++++++++++++++++-
 src/validation.cpp  |  4 ++-
 5 files changed, 96 insertions(+), 11 deletions(-)

diff --git a/src/core_write.cpp b/src/core_write.cpp
index 1771ebe0de..d9fc2a7a59 100644
--- a/src/core_write.cpp
+++ b/src/core_write.cpp
@@ -9,6 +9,7 @@
 #include <consensus/validation.h>
 #include <issuance.h>
 #include <key_io.h>
+#include <policy/discount.h> // ELEMENTS
 #include <script/descriptor.h>
 #include <script/script.h>
 #include <script/sign.h>
@@ -236,6 +237,10 @@ void TxToUniv(const CTransaction& tx, const uint256& hashBlock, UniValue& entry,
     entry.pushKV("version", static_cast<int64_t>(static_cast<uint32_t>(tx.nVersion)));
     entry.pushKV("size", (int)::GetSerializeSize(tx, PROTOCOL_VERSION));
     entry.pushKV("vsize", (GetTransactionWeight(tx) + WITNESS_SCALE_FACTOR - 1) / WITNESS_SCALE_FACTOR);
+    // ELEMENTS: add discountvsize
+    if (Params().GetAcceptDiscountCT()) {
+        entry.pushKV("discountvsize", GetDiscountVirtualTransactionSize(tx));
+    }
     entry.pushKV("weight", GetTransactionWeight(tx));
     entry.pushKV("locktime", (int64_t)tx.nLockTime);
 
diff --git a/src/policy/policy.h b/src/policy/policy.h
index f7e34546ed..f62f9a8640 100644
--- a/src/policy/policy.h
+++ b/src/policy/policy.h
@@ -7,6 +7,7 @@
 #define BITCOIN_POLICY_POLICY_H
 
 #include <consensus/consensus.h>
+#include <policy/discount.h>
 #include <policy/feerate.h>
 #include <script/interpreter.h>
 #include <script/standard.h>
diff --git a/src/txmempool.cpp b/src/txmempool.cpp
index c46b685b08..f84ffe4d23 100644
--- a/src/txmempool.cpp
+++ b/src/txmempool.cpp
@@ -13,6 +13,7 @@
 #include <consensus/tx_verify.h>
 #include <consensus/validation.h>
 #include <pegins.h>
+#include <policy/discount.h>
 #include <policy/fees.h>
 #include <policy/policy.h>
 #include <policy/settings.h>
@@ -43,18 +44,20 @@ struct update_descendant_state
 
 struct update_ancestor_state
 {
-    update_ancestor_state(int64_t _modifySize, CAmount _modifyFee, int64_t _modifyCount, int64_t _modifySigOpsCost) :
-        modifySize(_modifySize), modifyFee(_modifyFee), modifyCount(_modifyCount), modifySigOpsCost(_modifySigOpsCost)
+    update_ancestor_state(int64_t _modifySize, CAmount _modifyFee, int64_t _modifyCount, int64_t _modifySigOpsCost, int64_t _discountSize) :
+        modifySize(_modifySize), modifyFee(_modifyFee), modifyCount(_modifyCount), modifySigOpsCost(_modifySigOpsCost),
+        discountSize(_discountSize)
     {}
 
     void operator() (CTxMemPoolEntry &e)
-        { e.UpdateAncestorState(modifySize, modifyFee, modifyCount, modifySigOpsCost); }
+        { e.UpdateAncestorState(modifySize, modifyFee, modifyCount, modifySigOpsCost, discountSize); }
 
     private:
         int64_t modifySize;
         CAmount modifyFee;
         int64_t modifyCount;
         int64_t modifySigOpsCost;
+        int64_t discountSize;
 };
 
 struct update_fee_delta
@@ -102,6 +105,7 @@ CTxMemPoolEntry::CTxMemPoolEntry(const CTransactionRef& tx, CAmount fee,
       nSizeWithAncestors{GetTxSize()},
       nModFeesWithAncestors{nFee},
       nSigOpCostWithAncestors{sigOpCost},
+      discountSizeWithAncestors{GetDiscountTxSize()},
       setPeginsSpent(_setPeginsSpent) {}
 
 void CTxMemPoolEntry::UpdateFeeDelta(int64_t newFeeDelta)
@@ -121,6 +125,16 @@ size_t CTxMemPoolEntry::GetTxSize() const
     return GetVirtualTransactionSize(nTxWeight, sigOpCost);
 }
 
+size_t CTxMemPoolEntry::GetDiscountTxSize() const
+{
+    // discountvsize only differs from vsize if we accept discounted CTs
+    if (Params().GetAcceptDiscountCT()) {
+        return GetDiscountVirtualTransactionSize(*tx, sigOpCost, ::nBytesPerSigOp);
+    } else {
+        return GetVirtualTransactionSize(nTxWeight, sigOpCost);
+    }
+}
+
 void CTxMemPool::UpdateForDescendants(txiter updateIt, cacheMap& cachedDescendants,
                                       const std::set<uint256>& setExclude, std::set<uint256>& descendants_to_remove,
                                       uint64_t ancestor_size_limit, uint64_t ancestor_count_limit)
@@ -159,7 +173,7 @@ void CTxMemPool::UpdateForDescendants(txiter updateIt, cacheMap& cachedDescendan
             modifyCount++;
             cachedDescendants[updateIt].insert(mapTx.iterator_to(descendant));
             // Update ancestor state for each descendant
-            mapTx.modify(mapTx.iterator_to(descendant), update_ancestor_state(updateIt->GetTxSize(), updateIt->GetModifiedFee(), 1, updateIt->GetSigOpCost()));
+            mapTx.modify(mapTx.iterator_to(descendant), update_ancestor_state(updateIt->GetTxSize(), updateIt->GetModifiedFee(), 1, updateIt->GetSigOpCost(), updateIt->GetDiscountTxSize()));
             // Don't directly remove the transaction here -- doing so would
             // invalidate iterators in cachedDescendants. Mark it for removal
             // by inserting into descendants_to_remove.
@@ -370,12 +384,14 @@ void CTxMemPool::UpdateEntryForAncestors(txiter it, const setEntries &setAncesto
     int64_t updateSize = 0;
     CAmount updateFee = 0;
     int64_t updateSigOpsCost = 0;
+    int64_t discountSize = 0;
     for (txiter ancestorIt : setAncestors) {
         updateSize += ancestorIt->GetTxSize();
         updateFee += ancestorIt->GetModifiedFee();
         updateSigOpsCost += ancestorIt->GetSigOpCost();
+        discountSize += ancestorIt->GetDiscountTxSize();
     }
-    mapTx.modify(it, update_ancestor_state(updateSize, updateFee, updateCount, updateSigOpsCost));
+    mapTx.modify(it, update_ancestor_state(updateSize, updateFee, updateCount, updateSigOpsCost, discountSize));
 }
 
 void CTxMemPool::UpdateChildrenForRemoval(txiter it)
@@ -405,8 +421,9 @@ void CTxMemPool::UpdateForRemoveFromMempool(const setEntries &entriesToRemove, b
             int64_t modifySize = -((int64_t)removeIt->GetTxSize());
             CAmount modifyFee = -removeIt->GetModifiedFee();
             int modifySigOps = -removeIt->GetSigOpCost();
+            int64_t discountSize = -((int64_t)removeIt->GetDiscountTxSize());
             for (txiter dit : setDescendants) {
-                mapTx.modify(dit, update_ancestor_state(modifySize, modifyFee, -1, modifySigOps));
+                mapTx.modify(dit, update_ancestor_state(modifySize, modifyFee, -1, modifySigOps, discountSize));
             }
         }
     }
@@ -455,7 +472,7 @@ void CTxMemPoolEntry::UpdateDescendantState(int64_t modifySize, CAmount modifyFe
     assert(int64_t(nCountWithDescendants) > 0);
 }
 
-void CTxMemPoolEntry::UpdateAncestorState(int64_t modifySize, CAmount modifyFee, int64_t modifyCount, int64_t modifySigOps)
+void CTxMemPoolEntry::UpdateAncestorState(int64_t modifySize, CAmount modifyFee, int64_t modifyCount, int64_t modifySigOps, int64_t discountSize)
 {
     nSizeWithAncestors += modifySize;
     assert(int64_t(nSizeWithAncestors) > 0);
@@ -464,6 +481,8 @@ void CTxMemPoolEntry::UpdateAncestorState(int64_t modifySize, CAmount modifyFee,
     assert(int64_t(nCountWithAncestors) > 0);
     nSigOpCostWithAncestors += modifySigOps;
     assert(int(nSigOpCostWithAncestors) >= 0);
+    discountSizeWithAncestors += discountSize;
+    assert(int64_t(discountSizeWithAncestors) > 0);
 }
 
 CTxMemPool::CTxMemPool(CBlockPolicyEstimator* estimator, int check_ratio)
@@ -965,7 +984,7 @@ void CTxMemPool::queryHashes(std::vector<uint256>& vtxid) const
 }
 
 static TxMempoolInfo GetInfo(CTxMemPool::indexed_transaction_set::const_iterator it) {
-    return TxMempoolInfo{it->GetSharedTx(), it->GetTime(), it->GetFee(), it->GetTxSize(), it->GetModifiedFee() - it->GetFee()};
+    return TxMempoolInfo{it->GetSharedTx(), it->GetTime(), it->GetFee(), it->GetTxSize(), it->GetModifiedFee() - it->GetFee(), it->GetDiscountTxSize()};
 }
 
 std::vector<TxMempoolInfo> CTxMemPool::infoAll() const
@@ -1022,7 +1041,7 @@ void CTxMemPool::PrioritiseTransaction(const uint256& hash, const CAmount& nFeeD
             CalculateDescendants(it, setDescendants);
             setDescendants.erase(it);
             for (txiter descendantIt : setDescendants) {
-                mapTx.modify(descendantIt, update_ancestor_state(0, nFeeDelta, 0, 0));
+                mapTx.modify(descendantIt, update_ancestor_state(0, nFeeDelta, 0, 0, 0));
             }
             ++nTransactionsUpdated;
         }
diff --git a/src/txmempool.h b/src/txmempool.h
index 8b5e442a14..edfe7be1ad 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -117,6 +117,7 @@ private:
     uint64_t nSizeWithAncestors;
     CAmount nModFeesWithAncestors;
     int64_t nSigOpCostWithAncestors;
+    uint64_t discountSizeWithAncestors; // ELEMENTS
 
 public:
     CTxMemPoolEntry(const CTransactionRef& tx, CAmount fee,
@@ -129,6 +130,7 @@ public:
     CTransactionRef GetSharedTx() const { return this->tx; }
     const CAmount& GetFee() const { return nFee; }
     size_t GetTxSize() const;
+    size_t GetDiscountTxSize() const;
     size_t GetTxWeight() const { return nTxWeight; }
     std::chrono::seconds GetTime() const { return std::chrono::seconds{nTime}; }
     unsigned int GetHeight() const { return entryHeight; }
@@ -140,7 +142,7 @@ public:
     // Adjusts the descendant state.
     void UpdateDescendantState(int64_t modifySize, CAmount modifyFee, int64_t modifyCount);
     // Adjusts the ancestor state
-    void UpdateAncestorState(int64_t modifySize, CAmount modifyFee, int64_t modifyCount, int64_t modifySigOps);
+    void UpdateAncestorState(int64_t modifySize, CAmount modifyFee, int64_t modifyCount, int64_t modifySigOps, int64_t discountSize);
     // Updates the fee delta used for mining priority score, and the
     // modified fees with descendants.
     void UpdateFeeDelta(int64_t feeDelta);
@@ -155,6 +157,7 @@ public:
 
     uint64_t GetCountWithAncestors() const { return nCountWithAncestors; }
     uint64_t GetSizeWithAncestors() const { return nSizeWithAncestors; }
+    uint64_t GetDiscountSizeWithAncestors() const { return discountSizeWithAncestors; }
     CAmount GetModFeesWithAncestors() const { return nModFeesWithAncestors; }
     int64_t GetSigOpCostWithAncestors() const { return nSigOpCostWithAncestors; }
 
@@ -318,11 +321,56 @@ public:
     }
 };
 
+/** \class CompareTxMemPoolEntryByConfidentialFee
+ *
+ *  Sort an entry by min(score/discountvsize of entry's tx, score/discountvsize with all ancestors).
+ */
+class CompareTxMemPoolEntryByConfidentialFee
+{
+public:
+    template <typename T>
+    bool operator()(const T& a, const T& b) const
+    {
+        double a_mod_fee, a_size, b_mod_fee, b_size;
+
+        GetModFeeAndSize(a, a_mod_fee, a_size);
+        GetModFeeAndSize(b, b_mod_fee, b_size);
+
+        // Avoid division by rewriting (a/b > c/d) as (a*d > c*b).
+        double f1 = a_mod_fee * b_size;
+        double f2 = a_size * b_mod_fee;
+
+        if (f1 == f2) {
+            return a.GetTx().GetHash() < b.GetTx().GetHash();
+        }
+        return f1 > f2;
+    }
+
+    // Return the fee/size we're using for sorting this entry.
+    template <typename T>
+    void GetModFeeAndSize(const T& a, double& mod_fee, double& size) const
+    {
+        // Compare feerate with ancestors to feerate of the transaction, and
+        // return the fee/size for the min.
+        double f1 = (double)a.GetModifiedFee() * a.GetDiscountSizeWithAncestors();
+        double f2 = (double)a.GetModFeesWithAncestors() * a.GetDiscountTxSize();
+
+        if (f1 > f2) {
+            mod_fee = a.GetModFeesWithAncestors();
+            size = a.GetDiscountSizeWithAncestors();
+        } else {
+            mod_fee = a.GetModifiedFee();
+            size = a.GetDiscountTxSize();
+        }
+    }
+};
+
 // Multi_index tag names
 struct descendant_score {};
 struct entry_time {};
 struct ancestor_score {};
 struct index_by_wtxid {};
+struct confidential_score {}; // ELEMENTS
 
 class CBlockPolicyEstimator;
 
@@ -345,6 +393,9 @@ struct TxMempoolInfo
 
     /** The fee delta. */
     int64_t nFeeDelta;
+
+    /** ELEMENTS: Discounted CT size. */
+    size_t discountvsize;
 };
 
 /** Reason why a transaction was removed from the mempool,
@@ -380,6 +431,7 @@ enum class MemPoolRemovalReason {
  * - descendant feerate [we use max(feerate of tx, feerate of tx with all descendants)]
  * - time in mempool
  * - ancestor feerate [we use min(feerate of tx, feerate of tx with all unconfirmed ancestors)]
+ * - ancestor feerate by discount vsize // ELEMENTS: "confidential_score"
  *
  * Note: the term "descendant" refers to in-mempool transactions that depend on
  * this one, while "ancestor" refers to in-mempool transactions that a given
@@ -489,6 +541,12 @@ public:
                 boost::multi_index::tag<ancestor_score>,
                 boost::multi_index::identity<CTxMemPoolEntry>,
                 CompareTxMemPoolEntryByAncestorFee
+            >,
+            // ELEMENTS: sorted by confidential first fee rate with ancestors
+            boost::multi_index::ordered_non_unique<
+                boost::multi_index::tag<confidential_score>,
+                boost::multi_index::identity<CTxMemPoolEntry>,
+                CompareTxMemPoolEntryByConfidentialFee
             >
         >
     > indexed_transaction_set;
diff --git a/src/validation.cpp b/src/validation.cpp
index 1bd2e18ba3..2a2873aa83 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -914,7 +914,9 @@ bool MemPoolAccept::PreChecks(ATMPArgs& args, Workspace& ws)
 
     // No transactions are allowed below minRelayTxFee except from disconnected
     // blocks
-    if (!bypass_limits && !CheckFeeRate(ws.m_vsize, ws.m_modified_fees, state)) return false;
+    // ELEMENTS: accept discounted fees for Confidential Transactions only, if enabled.
+    int64_t package_size = Params().GetAcceptDiscountCT() ? GetDiscountVirtualTransactionSize(tx) : ws.m_vsize;
+    if (!bypass_limits && !CheckFeeRate(package_size, ws.m_modified_fees, state)) return false;
 
     ws.m_iters_conflicting = m_pool.GetIterSet(ws.m_conflicts);
     // Calculate in-mempool ancestors, up to a limit.