refactor: consolidate balance calculations (#338)

Move the account balance updates happening within `check_transaction` into `ManagedAccount::update_balance` and reuse it in `ManagedWalletInfo::update_balance` which also already gets called in `check_transaction`.

Author: xdustinface

key-wallet/src/managed_account/mod.rs

@@ -265,7 +265,20 @@ impl ManagedAccount {
     }
 
     /// Update the account balance
-    pub fn update_balance(&mut self, spendable: u64, unconfirmed: u64, locked: u64) {
+    pub fn update_balance(&mut self) {
+        let mut spendable = 0;
+        let mut unconfirmed = 0;
+        let mut locked = 0;
+        for utxo in self.utxos.values() {
+            let value = utxo.txout.value;
+            if utxo.is_locked {
+                locked += value;
+            } else if utxo.is_confirmed {
+                spendable += value;
+            } else {
+                unconfirmed += value;
+            }
+        }
         self.balance = WalletBalance::new(spendable, unconfirmed, locked);
         self.metadata.last_used = Some(Self::current_timestamp());
     }

key-wallet/src/transaction_checking/wallet_checker.rs

@@ -235,24 +235,6 @@ impl WalletTransactionChecker for ManagedWalletInfo {
                             for input in &tx.input {
                                 account.utxos.remove(&input.previous_output);
                             }
-
-                            // Recalculate account balance from UTXOs
-                            let mut confirmed = 0u64;
-                            let mut unconfirmed = 0u64;
-                            let mut locked = 0u64;
-
-                            for utxo in account.utxos.values() {
-                                let value = utxo.txout.value;
-                                if utxo.is_locked {
-                                    locked += value;
-                                } else if utxo.is_confirmed {
-                                    confirmed += value;
-                                } else {
-                                    unconfirmed += value;
-                                }
-                            }
-
-                            let _ = account.update_balance(confirmed, unconfirmed, locked);
                         }
                         _ => {
                             // Skip UTXO ingestion for identity/provider accounts

key-wallet/src/wallet/balance.rs

@@ -3,6 +3,7 @@
 //! This module provides a wallet balance structure containing all available balances.
 
 use core::fmt::{Display, Formatter};
+use core::ops::AddAssign;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
@@ -62,6 +63,14 @@ impl Display for WalletBalance {
     }
 }
 
+impl AddAssign for WalletBalance {
+    fn add_assign(&mut self, other: Self) {
+        self.spendable += other.spendable;
+        self.unconfirmed += other.unconfirmed;
+        self.locked += other.locked;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -91,4 +100,20 @@ mod tests {
         let display = balance.to_string();
         assert_eq!(display, "Spendable: 1000, Unconfirmed: 500, Locked: 200, Total: 1700");
     }
+
+    #[test]
+    fn test_balance_add_assign() {
+        let mut balance = WalletBalance::new(1000, 500, 200);
+        let balance_add = WalletBalance::new(300, 100, 50);
+        // Test adding actual balances
+        balance += balance_add;
+        assert_eq!(balance.spendable(), 1300);
+        assert_eq!(balance.unconfirmed(), 600);
+        assert_eq!(balance.locked(), 250);
+        assert_eq!(balance.total(), 2150);
+        // Test adding zero balances
+        let balance_before = balance;
+        balance += WalletBalance::default();
+        assert_eq!(balance_before, balance);
+    }
 }

key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs

@@ -3,6 +3,7 @@
 //! This trait allows WalletManager to work with different wallet info implementations
 
 use super::managed_account_operations::ManagedAccountOperations;
+use crate::account::ManagedAccountTrait;
 use crate::managed_account::managed_account_collection::ManagedAccountCollection;
 use crate::transaction_checking::WalletTransactionChecker;
 use crate::wallet::immature_transaction::{ImmatureTransaction, ImmatureTransactionCollection};
@@ -202,24 +203,12 @@ impl WalletInfoInterface for ManagedWalletInfo {
     }
 
     fn update_balance(&mut self) {
-        let mut spendable = 0u64;
-        let mut unconfirmed = 0u64;
-        let mut locked = 0u64;
-
-        for account in self.accounts.all_accounts() {
-            for utxo in account.utxos.values() {
-                let value = utxo.txout.value;
-                if utxo.is_locked {
-                    locked += value;
-                } else if utxo.is_confirmed {
-                    spendable += value;
-                } else {
-                    unconfirmed += value;
-                }
-            }
+        let mut balance = WalletBalance::default();
+        for account in self.accounts.all_accounts_mut() {
+            account.update_balance();
+            balance += *account.balance();
         }
-
-        self.balance = WalletBalance::new(spendable, unconfirmed, locked)
+        self.balance = balance;
     }
 
     fn transaction_history(&self) -> Vec<&TransactionRecord> {