wallet: Change balance calculation to use m_txos

Since we track the outputs owned by the wallet with m_txos, we can now
calculate the balance of the wallet by iterating m_txos and summing up
the amounts of the unspent txos.

As ISMINE_USED is not an actual isminetype that we attach to outputs and
was just passed into `CachedTxGetAvailableCredit` for convenience, we
pull the same determining logic from that function into `GetBalances` in
order to preserve existing behavior.

Author: achow101

src/wallet/receive.cpp

@@ -4,6 +4,7 @@
 
 #include <consensus/amount.h>
 #include <consensus/consensus.h>
+#include <util/check.h>
 #include <wallet/receive.h>
 #include <wallet/transaction.h>
 #include <wallet/wallet.h>
@@ -297,27 +298,41 @@ bool CachedTxIsTrusted(const CWallet& wallet, const CWalletTx& wtx)
 Balance GetBalance(const CWallet& wallet, const int min_depth, bool avoid_reuse)
 {
     Balance ret;
-    isminefilter reuse_filter = avoid_reuse ? ISMINE_NO : ISMINE_USED;
+    bool allow_used_addresses = !avoid_reuse || !wallet.IsWalletFlagSet(WALLET_FLAG_AVOID_REUSE);
     {
         LOCK(wallet.cs_wallet);
         std::set<Txid> trusted_parents;
-        for (const auto& entry : wallet.mapWallet)
-        {
-            const CWalletTx& wtx = entry.second;
+        for (const auto& [outpoint, txo] : wallet.GetTXOs()) {
+            const CWalletTx& wtx = txo.GetWalletTx();
+
             const bool is_trusted{CachedTxIsTrusted(wallet, wtx, trusted_parents)};
             const int tx_depth{wallet.GetTxDepthInMainChain(wtx)};
-            const CAmount tx_credit_mine{CachedTxGetAvailableCredit(wallet, wtx, ISMINE_SPENDABLE | reuse_filter)};
-            const CAmount tx_credit_watchonly{CachedTxGetAvailableCredit(wallet, wtx, ISMINE_WATCH_ONLY | reuse_filter)};
-            if (is_trusted && tx_depth >= min_depth) {
-                ret.m_mine_trusted += tx_credit_mine;
-                ret.m_watchonly_trusted += tx_credit_watchonly;
-            }
-            if (!is_trusted && tx_depth == 0 && wtx.InMempool()) {
-                ret.m_mine_untrusted_pending += tx_credit_mine;
-                ret.m_watchonly_untrusted_pending += tx_credit_watchonly;
+
+            if (!wallet.IsSpent(outpoint) && (allow_used_addresses || !wallet.IsSpentKey(txo.GetTxOut().scriptPubKey))) {
+                // Get the amounts for mine and watchonly
+                CAmount credit_mine = 0;
+                CAmount credit_watchonly = 0;
+                if (txo.GetIsMine() == ISMINE_SPENDABLE) {
+                    credit_mine = txo.GetTxOut().nValue;
+                } else if (txo.GetIsMine() == ISMINE_WATCH_ONLY) {
+                    credit_watchonly = txo.GetTxOut().nValue;
+                } else {
+                    // We shouldn't see any other isminetypes
+                    Assume(false);
+                }
+
+                // Set the amounts in the return object
+                if (wallet.IsTxImmatureCoinBase(wtx) && wtx.isConfirmed()) {
+                    ret.m_mine_immature += credit_mine;
+                    ret.m_watchonly_immature += credit_watchonly;
+                } else if (is_trusted && tx_depth >= min_depth) {
+                    ret.m_mine_trusted += credit_mine;
+                    ret.m_watchonly_trusted += credit_watchonly;
+                } else if (!is_trusted && wtx.InMempool()) {
+                    ret.m_mine_untrusted_pending += credit_mine;
+                    ret.m_watchonly_untrusted_pending += credit_watchonly;
+                }
             }
-            ret.m_mine_immature += CachedTxGetImmatureCredit(wallet, wtx, ISMINE_SPENDABLE);
-            ret.m_watchonly_immature += CachedTxGetImmatureCredit(wallet, wtx, ISMINE_WATCH_ONLY);
         }
     }
     return ret;
@@ -330,31 +345,21 @@ std::map<CTxDestination, CAmount> GetAddressBalances(const CWallet& wallet)
     {
         LOCK(wallet.cs_wallet);
         std::set<Txid> trusted_parents;
-        for (const auto& walletEntry : wallet.mapWallet)
-        {
-            const CWalletTx& wtx = walletEntry.second;
-
-            if (!CachedTxIsTrusted(wallet, wtx, trusted_parents))
-                continue;
+        for (const auto& [outpoint, txo] : wallet.GetTXOs()) {
+            const CWalletTx& wtx = txo.GetWalletTx();
 
-            if (wallet.IsTxImmatureCoinBase(wtx))
-                continue;
+            if (!CachedTxIsTrusted(wallet, wtx, trusted_parents)) continue;
+            if (wallet.IsTxImmatureCoinBase(wtx)) continue;
 
             int nDepth = wallet.GetTxDepthInMainChain(wtx);
-            if (nDepth < (CachedTxIsFromMe(wallet, wtx, ISMINE_ALL) ? 0 : 1))
-                continue;
+            if (nDepth < (CachedTxIsFromMe(wallet, wtx, ISMINE_ALL) ? 0 : 1)) continue;
 
-            for (unsigned int i = 0; i < wtx.tx->vout.size(); i++) {
-                const auto& output = wtx.tx->vout[i];
-                CTxDestination addr;
-                if (!wallet.IsMine(output))
-                    continue;
-                if(!ExtractDestination(output.scriptPubKey, addr))
-                    continue;
+            CTxDestination addr;
+            Assume(wallet.IsMine(txo.GetTxOut()));
+            if(!ExtractDestination(txo.GetTxOut().scriptPubKey, addr)) continue;
 
-                CAmount n = wallet.IsSpent(COutPoint(walletEntry.first, i)) ? 0 : output.nValue;
-                balances[addr] += n;
-            }
+            CAmount n = wallet.IsSpent(outpoint) ? 0 : txo.GetTxOut().nValue;
+            balances[addr] += n;
         }
     }
 

src/wallet/wallet.h

@@ -523,6 +523,8 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
 
     std::set<Txid> GetTxConflicts(const CWalletTx& wtx) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
+    const std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>& GetTXOs() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet) { AssertLockHeld(cs_wallet); return m_txos; };
+
     /** Cache outputs that belong to the wallet from a single transaction */
     void RefreshTXOsFromTx(const CWalletTx& wtx) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     /** Cache outputs that belong to the wallet for all transactions in the wallet */