wallet: single loop to calculate the currently required balances.

future: add watch-only and CS balances distinction when/if they are needed (at the moment, they aren't). Inspired in btc@fa57411fcba00556ba25d45bca53cc04623da051 .
PIVX-Project · Mar 11, 2021 · e9c160a · e9c160a
1 parent 32538ae
commit e9c160a
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 5 deletions.
diff --git a/src/interfaces/wallet.cpp b/src/interfaces/wallet.cpp
@@ -10,9 +10,10 @@ namespace interfaces {
 
     WalletBalances Wallet::getBalances() {
         WalletBalances result;
-        result.balance = m_wallet.GetAvailableBalance();
-        result.unconfirmed_balance = m_wallet.GetUnconfirmedBalance(ISMINE_SPENDABLE_TRANSPARENT);
-        result.immature_balance = m_wallet.GetImmatureBalance();
+        CWallet::Balance balance = m_wallet.GetBalance();
+        result.balance = balance.m_mine_trusted + balance.m_mine_trusted_shield;
+        result.unconfirmed_balance = balance.m_mine_untrusted_pending;
+        result.immature_balance = balance.m_mine_immature;
         result.have_watch_only = m_wallet.HaveWatchOnly();
         if (result.have_watch_only) {
             result.watch_only_balance = m_wallet.GetWatchOnlyBalance();
@@ -23,8 +24,8 @@ namespace interfaces {
             result.delegate_balance = m_wallet.GetDelegatedBalance();
             result.coldstaked_balance = m_wallet.GetColdStakingBalance();
         }
-        result.shielded_balance = m_wallet.GetAvailableShieldedBalance();
-        result.unconfirmed_shielded_balance = m_wallet.GetUnconfirmedShieldedBalance();
+        result.shielded_balance = balance.m_mine_trusted_shield;
+        result.unconfirmed_shielded_balance = balance.m_mine_untrusted_shielded_balance;
         return result;
     }
 

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
@@ -2055,6 +2055,32 @@ void CWallet::ResendWalletTransactions(CConnman* connman)
  * @{
  */
 
+CWallet::Balance CWallet::GetBalance(const int min_depth) const
+{
+    Balance ret;
+    {
+        LOCK(cs_wallet);
+        std::set<uint256> trusted_parents;
+        for (const auto& entry : mapWallet) {
+            const CWalletTx& wtx = entry.second;
+            const bool is_trusted{wtx.IsTrusted()};
+            const int tx_depth{wtx.GetDepthInMainChain()};
+            const CAmount tx_credit_mine{wtx.GetAvailableCredit(/* fUseCache */ true, ISMINE_SPENDABLE_TRANSPARENT)};
+            const CAmount tx_credit_shield_mine{wtx.GetAvailableCredit(/* fUseCache */ true, ISMINE_SPENDABLE_SHIELDED)};
+            if (is_trusted && tx_depth >= min_depth) {
+                ret.m_mine_trusted += tx_credit_mine;
+                ret.m_mine_trusted_shield += tx_credit_shield_mine;
+            }
+            if (!is_trusted && tx_depth == 0 && wtx.InMempool()) {
+                ret.m_mine_untrusted_pending += tx_credit_mine;
+                ret.m_mine_untrusted_shielded_balance += tx_credit_shield_mine;
+            }
+            ret.m_mine_immature += wtx.GetImmatureCredit();
+        }
+    }
+    return ret;
+}
+
 CAmount CWallet::loopTxsBalance(std::function<void(const uint256&, const CWalletTx&, CAmount&)> method) const
 {
     CAmount nTotal = 0;

diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -954,6 +954,15 @@ class CWallet : public CCryptoKeyStore, public CValidationInterface
     void ReacceptWalletTransactions(bool fFirstLoad = false);
     void ResendWalletTransactions(CConnman* connman) override;
 
+    struct Balance {
+        CAmount m_mine_trusted{0};               //!< Trusted, at depth=GetBalance.min_depth or more
+        CAmount m_mine_untrusted_pending{0};     //!< Untrusted, but in mempool (pending)
+        CAmount m_mine_immature{0};              //!< Immature coinbases/coinstakes in the main chain
+        CAmount m_mine_trusted_shield{0};        //!< Trusted shield, at depth=GetBalance.min_depth or more
+        CAmount m_mine_untrusted_shielded_balance{0}; //!< Untrusted shield, but in mempool (pending)
+    };
+    Balance GetBalance(int min_depth = 0) const;
+
     CAmount loopTxsBalance(std::function<void(const uint256&, const CWalletTx&, CAmount&)>method) const;
     CAmount GetAvailableBalance(bool fIncludeDelegated = true, bool fIncludeShielded = true) const;
     CAmount GetAvailableBalance(isminefilter& filter, bool useCache = false, int minDepth = 1) const;