From ae0876ec4273a8b978e2c602435a9ed25f48c976 Mon Sep 17 00:00:00 2001
From: Ava Chow <github@achow101.com>
Date: Tue, 20 Feb 2024 11:49:53 -0500
Subject: [PATCH] wallet: Keep track of transaction outputs owned by the wallet

When loading transactions to the wallet, check the outputs, and keep
track of the ones that are IsMine.
---
 src/wallet/transaction.h | 25 +++++++++++++++++++++++++
 src/wallet/wallet.cpp    | 26 ++++++++++++++++++++++++++
 src/wallet/wallet.h      |  6 ++++++
 3 files changed, 57 insertions(+)

diff --git a/src/wallet/transaction.h b/src/wallet/transaction.h
index 1ec62b992b6..8f64d181e0f 100644
--- a/src/wallet/transaction.h
+++ b/src/wallet/transaction.h
@@ -10,6 +10,7 @@
 #include <primitives/transaction.h>
 #include <tinyformat.h>
 #include <uint256.h>
+#include <util/check.h>
 #include <util/overloaded.h>
 #include <util/strencodings.h>
 #include <util/string.h>
@@ -362,6 +363,30 @@ struct WalletTxOrderComparator {
         return a->nOrderPos < b->nOrderPos;
     }
 };
+
+class WalletTXO
+{
+private:
+    const CWalletTx& m_wtx;
+    const CTxOut& m_output;
+    isminetype m_ismine;
+
+public:
+    WalletTXO(const CWalletTx& wtx, const CTxOut& output, const isminetype ismine)
+    : m_wtx(wtx),
+    m_output(output),
+    m_ismine(ismine)
+    {
+        Assume(std::ranges::find(wtx.tx->vout, output) != wtx.tx->vout.end());
+    }
+
+    const CWalletTx& GetWalletTx() const { return m_wtx; }
+
+    const CTxOut& GetTxOut() const { return m_output; }
+
+    isminetype GetIsMine() const { return m_ismine; }
+    void SetIsMine(isminetype ismine) { m_ismine = ismine; }
+};
 } // namespace wallet
 
 #endif // BITCOIN_WALLET_TRANSACTION_H
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 840d7d5bc94..49f1cf22bc5 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -1094,6 +1094,9 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const TxState& state, const
     // Break debit/credit balance caches:
     wtx.MarkDirty();
 
+    // Cache the outputs that belong to the wallet
+    RefreshTXOsFromTx(wtx);
+
     // Notify UI of new or updated transaction
     NotifyTransactionChanged(hash, fInsertedNew ? CT_NEW : CT_UPDATED);
 
@@ -1157,6 +1160,8 @@ bool CWallet::LoadToWallet(const Txid& hash, const UpdateWalletTxFn& fill_wtx)
     // Update birth time when tx time is older than it.
     MaybeUpdateBirthTime(wtx.GetTxTime());
 
+    // Make sure the tx outputs are known by the wallet
+    RefreshTXOsFromTx(wtx);
     return true;
 }
 
@@ -2324,6 +2329,9 @@ util::Result<void> CWallet::RemoveTxs(WalletBatch& batch, std::vector<Txid>& txs
             wtxOrdered.erase(it->second.m_it_wtxOrdered);
             for (const auto& txin : it->second.tx->vin)
                 mapTxSpends.erase(txin.prevout);
+            for (unsigned int i = 0; i < it->second.tx->vout.size(); ++i) {
+                m_txos.erase(COutPoint(Txid::FromUint256(hash), i));
+            }
             mapWallet.erase(it);
             NotifyTransactionChanged(hash, CT_DELETED);
         }
@@ -4425,4 +4433,22 @@ void CWallet::WriteBestBlock() const
         batch.WriteBestBlock(loc);
     }
 }
+
+void CWallet::RefreshTXOsFromTx(const CWalletTx& wtx)
+{
+    AssertLockHeld(cs_wallet);
+    for (uint32_t i = 0; i < wtx.tx->vout.size(); ++i) {
+        const CTxOut& txout = wtx.tx->vout.at(i);
+        isminetype ismine = IsMine(txout);
+        if (ismine == ISMINE_NO) {
+            continue;
+        }
+        COutPoint outpoint(wtx.GetHash(), i);
+        if (m_txos.contains(outpoint)) {
+            m_txos.at(outpoint).SetIsMine(ismine);
+        } else {
+            m_txos.emplace(outpoint, WalletTXO{wtx, txout, ismine});
+        }
+    }
+}
 } // namespace wallet
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 03c2ca2ffbf..4d7a369baf2 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -438,6 +438,9 @@ private:
     //! Cache of descriptor ScriptPubKeys used for IsMine. Maps ScriptPubKey to set of spkms
     std::unordered_map<CScript, std::vector<ScriptPubKeyMan*>, SaltedSipHasher> m_cached_spks;
 
+    //! Set of both spent and unspent transaction outputs owned by this wallet
+    std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher> m_txos GUARDED_BY(cs_wallet);
+
     /**
      * Catch wallet up to current chain, scanning new blocks, updating the best
      * block locator and m_last_block_processed, and registering for
@@ -520,6 +523,9 @@ public:
 
     std::set<Txid> GetTxConflicts(const CWalletTx& wtx) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
+    /** Cache outputs that belong to the wallet from a single transaction */
+    void RefreshTXOsFromTx(const CWalletTx& wtx) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+
     /**
      * Return depth of transaction in blockchain:
      * <0  : conflicts with a transaction this deep in the blockchain