From dde7cbe105ba6daaa636466d0fa3a83d15609417 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:49:58 -0500 Subject: [PATCH] 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. --- src/wallet/receive.cpp | 75 ++++++++++++++++++++++-------------------- src/wallet/wallet.h | 2 ++ 2 files changed, 42 insertions(+), 35 deletions(-) diff --git a/src/wallet/receive.cpp b/src/wallet/receive.cpp index 620d2bb12c6..30644b4becd 100644 --- a/src/wallet/receive.cpp +++ b/src/wallet/receive.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -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 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 (!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; + } } - if (!is_trusted && tx_depth == 0 && wtx.InMempool()) { - ret.m_mine_untrusted_pending += tx_credit_mine; - ret.m_watchonly_untrusted_pending += tx_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 GetAddressBalances(const CWallet& wallet) { LOCK(wallet.cs_wallet); std::set trusted_parents; - for (const auto& walletEntry : wallet.mapWallet) - { - const CWalletTx& wtx = walletEntry.second; + for (const auto& [outpoint, txo] : wallet.GetTXOs()) { + const CWalletTx& wtx = txo.GetWalletTx(); - if (!CachedTxIsTrusted(wallet, wtx, trusted_parents)) - continue; - - 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; } } diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index 2b52842236e..3d6f0508869 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -523,6 +523,8 @@ public: std::set GetTxConflicts(const CWalletTx& wtx) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet); + const std::unordered_map& 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 */