1 files changed, 45 insertions, 10 deletions

diff --git a/src/txmempool.h b/src/txmempool.h
index 583f7614b7..d4e9845942 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -198,6 +198,22 @@ struct mempoolentry_txid
     }
 };
 
+// extracts a transaction witness-hash from CTxMemPoolEntry or CTransactionRef
+struct mempoolentry_wtxid
+{
+    typedef uint256 result_type;
+    result_type operator() (const CTxMemPoolEntry &entry) const
+    {
+        return entry.GetTx().GetWitnessHash();
+    }
+
+    result_type operator() (const CTransactionRef& tx) const
+    {
+        return tx->GetWitnessHash();
+    }
+};
+
+
 /** \class CompareTxMemPoolEntryByDescendantScore
  *
  *  Sort an entry by max(score/size of entry's tx, score/size with all descendants).
@@ -318,6 +334,7 @@ public:
 struct descendant_score {};
 struct entry_time {};
 struct ancestor_score {};
+struct index_by_wtxid {};
 
 class CBlockPolicyEstimator;
 
@@ -383,8 +400,9 @@ public:
  *
  * CTxMemPool::mapTx, and CTxMemPoolEntry bookkeeping:
  *
- * mapTx is a boost::multi_index that sorts the mempool on 4 criteria:
- * - transaction hash
+ * mapTx is a boost::multi_index that sorts the mempool on 5 criteria:
+ * - transaction hash (txid)
+ * - witness-transaction hash (wtxid)
  * - 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)]
@@ -469,6 +487,12 @@ public:
         boost::multi_index::indexed_by<
             // sorted by txid
             boost::multi_index::hashed_unique<mempoolentry_txid, SaltedTxidHasher>,
+            // sorted by wtxid
+            boost::multi_index::hashed_unique<
+                boost::multi_index::tag<index_by_wtxid>,
+                mempoolentry_wtxid,
+                SaltedTxidHasher
+            >,
             // sorted by fee rate
             boost::multi_index::ordered_non_unique<
                 boost::multi_index::tag<descendant_score>,
@@ -549,8 +573,11 @@ private:
 
     std::vector<indexed_transaction_set::const_iterator> GetSortedDepthAndScore() const EXCLUSIVE_LOCKS_REQUIRED(cs);
 
-    /** track locally submitted transactions to periodically retry initial broadcast */
-    std::set<uint256> m_unbroadcast_txids GUARDED_BY(cs);
+    /**
+     * track locally submitted transactions to periodically retry initial broadcast
+     * map of txid -> wtxid
+     */
+    std::map<uint256, uint256> m_unbroadcast_txids GUARDED_BY(cs);
 
 public:
     indirectmap<COutPoint, const CTransaction*> mapNextTx GUARDED_BY(cs);
@@ -586,7 +613,7 @@ public:
 
     void clear();
     void _clear() EXCLUSIVE_LOCKS_REQUIRED(cs); //lock free
-    bool CompareDepthAndScore(const uint256& hasha, const uint256& hashb);
+    bool CompareDepthAndScore(const uint256& hasha, const uint256& hashb, bool wtxid=false);
     void queryHashes(std::vector<uint256>& vtxid) const;
     bool isSpent(const COutPoint& outpoint) const;
     unsigned int GetTransactionsUpdated() const;
@@ -689,24 +716,32 @@ public:
         return totalTxSize;
     }
 
-    bool exists(const uint256& hash) const
+    bool exists(const uint256& hash, bool wtxid=false) const
     {
         LOCK(cs);
+        if (wtxid) {
+            return (mapTx.get<index_by_wtxid>().count(hash) != 0);
+        }
         return (mapTx.count(hash) != 0);
     }
 
     CTransactionRef get(const uint256& hash) const;
-    TxMempoolInfo info(const uint256& hash) const;
+    txiter get_iter_from_wtxid(const uint256& wtxid) const EXCLUSIVE_LOCKS_REQUIRED(cs)
+    {
+        AssertLockHeld(cs);
+        return mapTx.project<0>(mapTx.get<index_by_wtxid>().find(wtxid));
+    }
+    TxMempoolInfo info(const uint256& hash, bool wtxid=false) const;
     std::vector<TxMempoolInfo> infoAll() const;
 
     size_t DynamicMemoryUsage() const;
 
     /** Adds a transaction to the unbroadcast set */
-    void AddUnbroadcastTx(const uint256& txid) {
+    void AddUnbroadcastTx(const uint256& txid, const uint256& wtxid) {
         LOCK(cs);
         // Sanity Check: the transaction should also be in the mempool
         if (exists(txid)) {
-            m_unbroadcast_txids.insert(txid);
+            m_unbroadcast_txids[txid] = wtxid;
         }
     }
 
@@ -714,7 +749,7 @@ public:
     void RemoveUnbroadcastTx(const uint256& txid, const bool unchecked = false);
 
     /** Returns transactions in unbroadcast set */
-    std::set<uint256> GetUnbroadcastTxs() const {
+    std::map<uint256, uint256> GetUnbroadcastTxs() const {
         LOCK(cs);
         return m_unbroadcast_txids;
     }