diff --git a/key-wallet-ffi/FFI_API.md b/key-wallet-ffi/FFI_API.md index 7d4675863..b19cf3e67 100644 --- a/key-wallet-ffi/FFI_API.md +++ b/key-wallet-ffi/FFI_API.md @@ -915,14 +915,14 @@ Get address pool information for an account # Safety - `managed_wallet` must b #### `managed_wallet_get_balance` ```c -managed_wallet_get_balance(managed_wallet: *const FFIManagedWalletInfo, confirmed_out: *mut u64, unconfirmed_out: *mut u64, locked_out: *mut u64, total_out: *mut u64, error: *mut FFIError,) -> bool +managed_wallet_get_balance(managed_wallet: *const FFIManagedWalletInfo, confirmed_out: *mut u64, unconfirmed_out: *mut u64, immature_out: *mut u64, locked_out: *mut u64, total_out: *mut u64, error: *mut FFIError,) -> bool ``` **Description:** -Get wallet balance from managed wallet info Returns the balance breakdown including confirmed, unconfirmed, locked, and total amounts. # Safety - `managed_wallet` must be a valid pointer to an FFIManagedWalletInfo - `confirmed_out` must be a valid pointer to store the confirmed balance - `unconfirmed_out` must be a valid pointer to store the unconfirmed balance - `locked_out` must be a valid pointer to store the locked balance - `total_out` must be a valid pointer to store the total balance - `error` must be a valid pointer to an FFIError +Get wallet balance from managed wallet info Returns the balance breakdown including confirmed, unconfirmed, immature, locked, and total amounts. # Safety - `managed_wallet` must be a valid pointer to an FFIManagedWalletInfo - `confirmed_out` must be a valid pointer to store the confirmed balance - `unconfirmed_out` must be a valid pointer to store the unconfirmed balance - `immature_out` must be a valid pointer to store the immature balance - `locked_out` must be a valid pointer to store the locked balance - `total_out` must be a valid pointer to store the total balance - `error` must be a valid pointer to an FFIError **Safety:** -- `managed_wallet` must be a valid pointer to an FFIManagedWalletInfo - `confirmed_out` must be a valid pointer to store the confirmed balance - `unconfirmed_out` must be a valid pointer to store the unconfirmed balance - `locked_out` must be a valid pointer to store the locked balance - `total_out` must be a valid pointer to store the total balance - `error` must be a valid pointer to an FFIError +- `managed_wallet` must be a valid pointer to an FFIManagedWalletInfo - `confirmed_out` must be a valid pointer to store the confirmed balance - `unconfirmed_out` must be a valid pointer to store the unconfirmed balance - `immature_out` must be a valid pointer to store the immature balance - `locked_out` must be a valid pointer to store the locked balance - `total_out` must be a valid pointer to store the total balance - `error` must be a valid pointer to an FFIError **Module:** `managed_wallet` diff --git a/key-wallet-ffi/include/key_wallet_ffi.h b/key-wallet-ffi/include/key_wallet_ffi.h index 7da7c404c..210dea707 100644 --- a/key-wallet-ffi/include/key_wallet_ffi.h +++ b/key-wallet-ffi/include/key_wallet_ffi.h @@ -516,11 +516,15 @@ typedef struct { */ uint64_t unconfirmed; /* - Immature balance in duffs (e.g., mining rewards) + Immature balance in duffs (e.g., mining rewards not yet mature) */ uint64_t immature; /* - Total balance (confirmed + unconfirmed) in duffs + Locked balance in duffs (e.g., CoinJoin reserves) + */ + uint64_t locked; + /* + Total balance in duffs */ uint64_t total; } FFIBalance; @@ -3069,13 +3073,14 @@ bool managed_wallet_get_bip_44_internal_address_range(FFIManagedWalletInfo *mana /* Get wallet balance from managed wallet info - Returns the balance breakdown including confirmed, unconfirmed, locked, and total amounts. + Returns the balance breakdown including confirmed, unconfirmed, immature, locked, and total amounts. # Safety - `managed_wallet` must be a valid pointer to an FFIManagedWalletInfo - `confirmed_out` must be a valid pointer to store the confirmed balance - `unconfirmed_out` must be a valid pointer to store the unconfirmed balance + - `immature_out` must be a valid pointer to store the immature balance - `locked_out` must be a valid pointer to store the locked balance - `total_out` must be a valid pointer to store the total balance - `error` must be a valid pointer to an FFIError @@ -3084,6 +3089,7 @@ bool managed_wallet_get_bip_44_internal_address_range(FFIManagedWalletInfo *mana bool managed_wallet_get_balance(const FFIManagedWalletInfo *managed_wallet, uint64_t *confirmed_out, uint64_t *unconfirmed_out, + uint64_t *immature_out, uint64_t *locked_out, uint64_t *total_out, FFIError *error) diff --git a/key-wallet-ffi/src/managed_account.rs b/key-wallet-ffi/src/managed_account.rs index 10d7afee3..54c9bbbef 100644 --- a/key-wallet-ffi/src/managed_account.rs +++ b/key-wallet-ffi/src/managed_account.rs @@ -515,7 +515,8 @@ pub unsafe extern "C" fn managed_account_get_balance( *balance_out = crate::types::FFIBalance { confirmed: balance.spendable(), unconfirmed: balance.unconfirmed(), - immature: 0, // WalletBalance doesn't have immature field + immature: balance.immature(), + locked: balance.locked(), total: balance.total(), }; @@ -1236,6 +1237,7 @@ mod tests { confirmed: 999, unconfirmed: 999, immature: 999, + locked: 999, total: 999, }; let success = managed_account_get_balance(account, &mut balance_out); @@ -1244,6 +1246,7 @@ mod tests { assert_eq!(balance_out.confirmed, 0); assert_eq!(balance_out.unconfirmed, 0); assert_eq!(balance_out.immature, 0); + assert_eq!(balance_out.locked, 0); assert_eq!(balance_out.total, 0); // Test get_transaction_count diff --git a/key-wallet-ffi/src/managed_wallet.rs b/key-wallet-ffi/src/managed_wallet.rs index 81a3f111a..00b0bcf65 100644 --- a/key-wallet-ffi/src/managed_wallet.rs +++ b/key-wallet-ffi/src/managed_wallet.rs @@ -531,13 +531,14 @@ pub unsafe extern "C" fn managed_wallet_get_bip_44_internal_address_range( /// Get wallet balance from managed wallet info /// -/// Returns the balance breakdown including confirmed, unconfirmed, locked, and total amounts. +/// Returns the balance breakdown including confirmed, unconfirmed, immature, locked, and total amounts. /// /// # Safety /// /// - `managed_wallet` must be a valid pointer to an FFIManagedWalletInfo /// - `confirmed_out` must be a valid pointer to store the confirmed balance /// - `unconfirmed_out` must be a valid pointer to store the unconfirmed balance +/// - `immature_out` must be a valid pointer to store the immature balance /// - `locked_out` must be a valid pointer to store the locked balance /// - `total_out` must be a valid pointer to store the total balance /// - `error` must be a valid pointer to an FFIError @@ -546,6 +547,7 @@ pub unsafe extern "C" fn managed_wallet_get_balance( managed_wallet: *const FFIManagedWalletInfo, confirmed_out: *mut u64, unconfirmed_out: *mut u64, + immature_out: *mut u64, locked_out: *mut u64, total_out: *mut u64, error: *mut FFIError, @@ -561,6 +563,7 @@ pub unsafe extern "C" fn managed_wallet_get_balance( if confirmed_out.is_null() || unconfirmed_out.is_null() + || immature_out.is_null() || locked_out.is_null() || total_out.is_null() { @@ -578,6 +581,7 @@ pub unsafe extern "C" fn managed_wallet_get_balance( unsafe { *confirmed_out = balance.spendable(); *unconfirmed_out = balance.unconfirmed(); + *immature_out = balance.immature(); *locked_out = balance.locked(); *total_out = balance.total(); } @@ -1042,7 +1046,7 @@ mod tests { let mut managed_info = ManagedWalletInfo::from_wallet(wallet_arc); // Set some test balance values - managed_info.balance = WalletBalance::new(1000000, 50000, 25000); + managed_info.balance = WalletBalance::new(1000000, 50000, 10000, 25000); let ffi_managed = FFIManagedWalletInfo::new(managed_info); let ffi_managed_ptr = Box::into_raw(Box::new(ffi_managed)); @@ -1050,6 +1054,7 @@ mod tests { // Test getting balance let mut confirmed: u64 = 0; let mut unconfirmed: u64 = 0; + let mut immature: u64 = 0; let mut locked: u64 = 0; let mut total: u64 = 0; @@ -1058,6 +1063,7 @@ mod tests { ffi_managed_ptr, &mut confirmed, &mut unconfirmed, + &mut immature, &mut locked, &mut total, &mut error, @@ -1067,8 +1073,9 @@ mod tests { assert!(success); assert_eq!(confirmed, 1000000); assert_eq!(unconfirmed, 50000); + assert_eq!(immature, 10000); assert_eq!(locked, 25000); - assert_eq!(total, 1075000); + assert_eq!(total, 1085000); // Test with null managed wallet let success = unsafe { @@ -1076,6 +1083,7 @@ mod tests { ptr::null(), &mut confirmed, &mut unconfirmed, + &mut immature, &mut locked, &mut total, &mut error, @@ -1091,6 +1099,7 @@ mod tests { ffi_managed_ptr, ptr::null_mut(), &mut unconfirmed, + &mut immature, &mut locked, &mut total, &mut error, diff --git a/key-wallet-ffi/src/types.rs b/key-wallet-ffi/src/types.rs index f8c4217f2..8a8c94b47 100644 --- a/key-wallet-ffi/src/types.rs +++ b/key-wallet-ffi/src/types.rs @@ -55,9 +55,11 @@ pub struct FFIBalance { pub confirmed: u64, /// Unconfirmed balance in duffs pub unconfirmed: u64, - /// Immature balance in duffs (e.g., mining rewards) + /// Immature balance in duffs (e.g., mining rewards not yet mature) pub immature: u64, - /// Total balance (confirmed + unconfirmed) in duffs + /// Locked balance in duffs (e.g., CoinJoin reserves) + pub locked: u64, + /// Total balance in duffs pub total: u64, } @@ -66,7 +68,8 @@ impl From for FFIBalance { FFIBalance { confirmed: balance.spendable(), unconfirmed: balance.unconfirmed(), - immature: balance.locked(), // Map locked to immature for now + immature: balance.immature(), + locked: balance.locked(), total: balance.total(), } } diff --git a/key-wallet-manager/tests/spv_integration_tests.rs b/key-wallet-manager/tests/spv_integration_tests.rs index 2201f5913..c285cf1c5 100644 --- a/key-wallet-manager/tests/spv_integration_tests.rs +++ b/key-wallet-manager/tests/spv_integration_tests.rs @@ -4,6 +4,7 @@ use dashcore::bip158::{BlockFilter, BlockFilterWriter}; use dashcore::blockdata::block::{Block, Header, Version}; use dashcore::blockdata::script::ScriptBuf; use dashcore::blockdata::transaction::Transaction; +use dashcore::constants::COINBASE_MATURITY; use dashcore::pow::CompactTarget; use dashcore::{BlockHash, OutPoint, TxIn, TxOut, Txid}; use dashcore_hashes::Hash; @@ -199,6 +200,29 @@ fn assert_wallet_heights(manager: &WalletManager, expected_he } } +/// Create a coinbase transaction paying to the given script +/// TODO: Unify with other `create_coinbase_transaction` helpers into `dashcore` crate. +fn create_coinbase_transaction(script_pubkey: ScriptBuf, value: u64) -> Transaction { + Transaction { + version: 2, + lock_time: 0, + input: vec![TxIn { + previous_output: OutPoint { + txid: Txid::all_zeros(), + vout: 0xffffffff, + }, + script_sig: ScriptBuf::new(), + sequence: 0xffffffff, + witness: dashcore::Witness::default(), + }], + output: vec![TxOut { + value, + script_pubkey, + }], + special_transaction_payload: None, + } +} + /// Test that the wallet heights are updated after block processing. #[tokio::test] async fn test_height_updated_after_block_processing() { @@ -218,3 +242,81 @@ async fn test_height_updated_after_block_processing() { assert_wallet_heights(&manager, height); } } + +#[tokio::test] +async fn test_immature_balance_matures_during_block_processing() { + let mut manager = WalletManager::::new(Network::Testnet); + + // Create a wallet and get an address to receive the coinbase + let wallet_id = manager + .create_wallet_with_random_mnemonic(WalletAccountCreationOptions::Default) + .expect("Failed to create wallet"); + + let account_xpub = { + let wallet = manager.get_wallet(&wallet_id).expect("Wallet should exist"); + wallet.accounts.standard_bip44_accounts.get(&0).expect("Should have account").account_xpub + }; + + // Get the first receive address from the wallet + let receive_address = { + let wallet_info = + manager.get_wallet_info_mut(&wallet_id).expect("Wallet info should exist"); + wallet_info + .first_bip44_managed_account_mut() + .expect("Should have managed account") + .next_receive_address(Some(&account_xpub), true) + .expect("Should get address") + }; + + // Create a coinbase transaction paying to our wallet + let coinbase_value = 100; + let coinbase_tx = create_coinbase_transaction(receive_address.script_pubkey(), coinbase_value); + + // Process the coinbase at height 1000 + let coinbase_height = 1000; + let coinbase_block = create_test_block(coinbase_height, vec![coinbase_tx.clone()]); + manager.process_block(&coinbase_block, coinbase_height).await; + + // Verify the coinbase is detected and stored as immature + let wallet_info = manager.get_wallet_info(&wallet_id).expect("Wallet info should exist"); + assert!( + wallet_info.immature_transactions().contains(&coinbase_tx), + "Coinbase should be in immature transactions" + ); + assert_eq!( + wallet_info.balance().immature(), + coinbase_value, + "Immature balance should reflect coinbase" + ); + + // Process 99 more blocks up to just before maturity + let maturity_height = coinbase_height + COINBASE_MATURITY; + for height in (coinbase_height + 1)..maturity_height { + let block = create_test_block(height, vec![create_test_transaction(1000)]); + manager.process_block(&block, height).await; + } + + // Verify still immature just before maturity + let wallet_info = manager.get_wallet_info(&wallet_id).expect("Wallet info should exist"); + assert!( + wallet_info.immature_transactions().contains(&coinbase_tx), + "Coinbase should still be immature at height {}", + maturity_height - 1 + ); + + // Process the maturity block + let maturity_block = create_test_block(maturity_height, vec![create_test_transaction(1000)]); + manager.process_block(&maturity_block, maturity_height).await; + + // Verify the coinbase has matured + let wallet_info = manager.get_wallet_info(&wallet_id).expect("Wallet info should exist"); + assert!( + !wallet_info.immature_transactions().contains(&coinbase_tx), + "Coinbase should no longer be immature after maturity height" + ); + assert_eq!( + wallet_info.balance().immature(), + 0, + "Immature balance should be zero after maturity" + ); +} diff --git a/key-wallet/src/managed_account/mod.rs b/key-wallet/src/managed_account/mod.rs index 5548ec650..14be9a322 100644 --- a/key-wallet/src/managed_account/mod.rs +++ b/key-wallet/src/managed_account/mod.rs @@ -268,18 +268,21 @@ impl ManagedAccount { pub fn update_balance(&mut self, synced_height: u32) { let mut spendable = 0; let mut unconfirmed = 0; + let mut immature = 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_mature(synced_height) { + immature += value; } else if utxo.is_spendable(synced_height) { spendable += value; } else { unconfirmed += value; } } - self.balance = WalletBalance::new(spendable, unconfirmed, locked); + self.balance = WalletBalance::new(spendable, unconfirmed, immature, locked); self.metadata.last_used = Some(Self::current_timestamp()); } diff --git a/key-wallet/src/test_utils/account.rs b/key-wallet/src/test_utils/account.rs new file mode 100644 index 000000000..bf3f427cc --- /dev/null +++ b/key-wallet/src/test_utils/account.rs @@ -0,0 +1,39 @@ +use crate::account::StandardAccountType; +use crate::managed_account::address_pool::{AddressPool, AddressPoolType, KeySource}; +use crate::managed_account::managed_account_type::ManagedAccountType; +use crate::managed_account::ManagedAccount; +use crate::{DerivationPath, Network}; + +impl ManagedAccount { + /// Create a test managed account with a standard BIP44 type and empty address pools + pub fn new_test_bip44(network: Network) -> Self { + let base_path = DerivationPath::master(); + + let external_pool = AddressPool::new( + base_path.clone(), + AddressPoolType::External, + 20, + network, + &KeySource::NoKeySource, + ) + .expect("Failed to create external address pool"); + + let internal_pool = AddressPool::new( + base_path, + AddressPoolType::Internal, + 20, + network, + &KeySource::NoKeySource, + ) + .expect("Failed to create internal address pool"); + + let account_type = ManagedAccountType::Standard { + index: 0, + standard_account_type: StandardAccountType::BIP44Account, + external_addresses: external_pool, + internal_addresses: internal_pool, + }; + + ManagedAccount::new(account_type, network, false) + } +} diff --git a/key-wallet/src/test_utils/mod.rs b/key-wallet/src/test_utils/mod.rs index ad2592879..baaa97e9d 100644 --- a/key-wallet/src/test_utils/mod.rs +++ b/key-wallet/src/test_utils/mod.rs @@ -1,3 +1,4 @@ +mod account; mod address; mod utxo; diff --git a/key-wallet/src/tests/balance_tests.rs b/key-wallet/src/tests/balance_tests.rs new file mode 100644 index 000000000..38fc78d88 --- /dev/null +++ b/key-wallet/src/tests/balance_tests.rs @@ -0,0 +1,81 @@ +//! Tests for update_balance() UTXO categorization. + +use crate::managed_account::ManagedAccount; +use crate::wallet::managed_wallet_info::wallet_info_interface::WalletInfoInterface; +use crate::wallet::managed_wallet_info::ManagedWalletInfo; +use crate::{Network, Utxo, WalletBalance}; + +#[test] +fn test_balance_with_mixed_utxo_types() { + let mut wallet_info = ManagedWalletInfo::new(Network::Regtest, [1u8; 32]); + let mut account = ManagedAccount::new_test_bip44(Network::Regtest); + + // Regular confirmed UTXO + let utxo1 = Utxo::new_test(1, 100_000, 1000, false, true); + account.utxos.insert(utxo1.outpoint, utxo1); + // Mature coinbase (100+ confirmations at height 1100) + let utxo2 = Utxo::new_test(2, 10_000_000, 1000, true, true); + account.utxos.insert(utxo2.outpoint, utxo2); + // Immature coinbase (<100 confirmations at height 1100) + let utxo3 = Utxo::new_test(3, 20_000_000, 1050, true, true); + account.utxos.insert(utxo3.outpoint, utxo3); + wallet_info.accounts.insert(account); + + assert_eq!(wallet_info.balance(), WalletBalance::default()); + wallet_info.update_synced_height(1100); + let expected = WalletBalance::new(10_100_000, 0, 20_000_000, 0); + assert_eq!(wallet_info.balance(), expected); +} + +#[test] +fn test_coinbase_maturity_boundary() { + let mut wallet_info = ManagedWalletInfo::new(Network::Regtest, [2u8; 32]); + let mut account = ManagedAccount::new_test_bip44(Network::Regtest); + + // Coinbase at height 1000 + let utxo = Utxo::new_test(1, 50_000_000, 1000, true, true); + account.utxos.insert(utxo.outpoint, utxo); + wallet_info.accounts.insert(account); + + assert_eq!(wallet_info.balance(), WalletBalance::default()); + // 99 confirmations: immature + wallet_info.update_synced_height(1099); + let expected_immature = WalletBalance::new(0, 0, 50_000_000, 0); + assert_eq!(wallet_info.balance(), expected_immature); + + // 100 confirmations: mature + wallet_info.update_synced_height(1100); + let expected_mature = WalletBalance::new(50_000_000, 0, 0, 0); + assert_eq!(wallet_info.balance(), expected_mature); +} + +#[test] +fn test_locked_utxos_in_locked_balance() { + let mut wallet_info = ManagedWalletInfo::new(Network::Regtest, [3u8; 32]); + let mut account = ManagedAccount::new_test_bip44(Network::Regtest); + + let mut utxo = Utxo::new_test(1, 100_000, 1000, false, true); + utxo.is_locked = true; + account.utxos.insert(utxo.outpoint, utxo); + wallet_info.accounts.insert(account); + + assert_eq!(wallet_info.balance(), WalletBalance::default()); + wallet_info.update_synced_height(1100); + let expected = WalletBalance::new(0, 0, 0, 100_000); + assert_eq!(wallet_info.balance(), expected); +} + +#[test] +fn test_unconfirmed_utxos_in_unconfirmed_balance() { + let mut wallet_info = ManagedWalletInfo::new(Network::Regtest, [4u8; 32]); + let mut account = ManagedAccount::new_test_bip44(Network::Regtest); + + let utxo = Utxo::new_test(1, 100_000, 0, false, false); + account.utxos.insert(utxo.outpoint, utxo); + wallet_info.accounts.insert(account); + + assert_eq!(wallet_info.balance(), WalletBalance::default()); + wallet_info.update_synced_height(1100); + let expected = WalletBalance::new(0, 100_000, 0, 0); + assert_eq!(wallet_info.balance(), expected); +} diff --git a/key-wallet/src/tests/mod.rs b/key-wallet/src/tests/mod.rs index 1e59f2e3e..339ff67db 100644 --- a/key-wallet/src/tests/mod.rs +++ b/key-wallet/src/tests/mod.rs @@ -4,6 +4,8 @@ mod account_tests; +mod balance_tests; + mod address_pool_tests; mod advanced_transaction_tests; diff --git a/key-wallet/src/transaction_checking/wallet_checker.rs b/key-wallet/src/transaction_checking/wallet_checker.rs index 2872fec18..04f7a8638 100644 --- a/key-wallet/src/transaction_checking/wallet_checker.rs +++ b/key-wallet/src/transaction_checking/wallet_checker.rs @@ -5,7 +5,6 @@ pub(crate) use super::account_checker::TransactionCheckResult; use super::transaction_router::TransactionRouter; -use crate::wallet::managed_wallet_info::wallet_info_interface::WalletInfoInterface; use crate::wallet::managed_wallet_info::ManagedWalletInfo; use crate::{Utxo, Wallet}; use async_trait::async_trait; @@ -289,9 +288,6 @@ impl WalletTransactionChecker for ManagedWalletInfo { // Update wallet metadata self.metadata.total_transactions += 1; - // Update cached balance - self.update_balance(); - // Log the detected transaction let wallet_net: i64 = (result.total_received as i64) - (result.total_sent as i64); let ctx = match context { @@ -545,8 +541,6 @@ mod tests { }; let block_height = 100000; - // Set synced_height to the block height where coinbase was received - managed_wallet.update_synced_height(block_height); // Test with InBlock context let context = TransactionContext::InBlock { @@ -557,6 +551,8 @@ mod tests { let result = managed_wallet.check_transaction(&coinbase_tx, context, &mut wallet, true).await; + // Set synced_height to block where coinbase was received to trigger balance updates. + managed_wallet.update_synced_height(block_height); // Should be relevant assert!(result.is_relevant); @@ -580,7 +576,7 @@ mod tests { assert_eq!(immature_txs[0].txid(), coinbase_tx.txid()); // Immature balance should reflect the coinbase value - assert_eq!(managed_wallet.immature_balance(), 5_000_000_000); + assert_eq!(managed_wallet.balance().immature(), 5_000_000_000); // Spendable UTXOs should be empty (coinbase not mature) assert!( @@ -720,8 +716,6 @@ mod tests { }; let block_height = 100000; - // Set synced_height to block where coinbase was received - managed_wallet.update_synced_height(block_height); let context = TransactionContext::InBlock { height: block_height, @@ -732,6 +726,8 @@ mod tests { // Process the coinbase transaction let result = managed_wallet.check_transaction(&coinbase_tx, context, &mut wallet, true).await; + // Set synced_height to block where coinbase was received to trigger balance updates. + managed_wallet.update_synced_height(block_height); // Should be relevant assert!(result.is_relevant); @@ -754,10 +750,12 @@ mod tests { assert_eq!(immature_txs.len(), 1, "Should have one immature transaction"); // Immature balance should reflect the coinbase value - let immature_balance = managed_wallet.immature_balance(); - assert_eq!( - immature_balance, 5_000_000_000, - "Immature balance should reflect the coinbase value" + assert_eq!(managed_wallet.balance().immature(), 5_000_000_000); + + // Spendable UTXOs should be empty (coinbase not mature yet) + assert!( + managed_wallet.get_spendable_utxos().is_empty(), + "No spendable UTXOs while coinbase is immature" ); // Spendable UTXOs should be empty (coinbase not mature yet) @@ -782,7 +780,7 @@ mod tests { assert!(immature_txs.is_empty(), "Matured coinbase should not be in immature transactions"); // Immature balance should now be zero - let immature_balance = managed_wallet.immature_balance(); + let immature_balance = managed_wallet.balance().immature(); assert_eq!(immature_balance, 0, "Immature balance should be zero after maturity"); // Spendable UTXOs should now contain the matured coinbase diff --git a/key-wallet/src/wallet/balance.rs b/key-wallet/src/wallet/balance.rs index 08467aa6c..55d402482 100644 --- a/key-wallet/src/wallet/balance.rs +++ b/key-wallet/src/wallet/balance.rs @@ -15,16 +15,19 @@ pub struct WalletBalance { spendable: u64, /// Unconfirmed balance (UTXOs without confirmations) unconfirmed: u64, + /// Immature balance (UTXOs without enough confirmations for maturity. e.g., 100 for coinbase.) + immature: u64, /// Locked balance (UTXOs reserved for specific purposes like CoinJoin) locked: u64, } impl WalletBalance { /// Create a new wallet balance - pub fn new(spendable: u64, unconfirmed: u64, locked: u64) -> Self { + pub fn new(spendable: u64, unconfirmed: u64, immature: u64, locked: u64) -> Self { Self { spendable, unconfirmed, + immature, locked, } } @@ -39,6 +42,11 @@ impl WalletBalance { self.unconfirmed } + /// Get the immature balance. + pub fn immature(&self) -> u64 { + self.immature + } + /// Get the locked balance. pub fn locked(&self) -> u64 { self.locked @@ -46,7 +54,7 @@ impl WalletBalance { /// Get the total balance. pub fn total(&self) -> u64 { - self.spendable + self.unconfirmed + self.locked + self.spendable + self.unconfirmed + self.immature + self.locked } } @@ -54,9 +62,10 @@ impl Display for WalletBalance { fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { write!( f, - "Spendable: {}, Unconfirmed: {}, Locked: {}, Total: {}", + "Spendable: {}, Unconfirmed: {}, Immature: {}, Locked: {}, Total: {}", self.spendable, self.unconfirmed, + self.immature, self.locked, self.total() ) @@ -67,6 +76,7 @@ impl AddAssign for WalletBalance { fn add_assign(&mut self, other: Self) { self.spendable += other.spendable; self.unconfirmed += other.unconfirmed; + self.immature += other.immature; self.locked += other.locked; } } @@ -77,40 +87,48 @@ mod tests { #[test] fn test_balance_creation_and_getters() { - let balance = WalletBalance::new(1000, 500, 200); + let balance = WalletBalance::new(1000, 500, 100, 200); assert_eq!(balance.spendable(), 1000); assert_eq!(balance.unconfirmed(), 500); + assert_eq!(balance.immature(), 100); assert_eq!(balance.locked(), 200); - assert_eq!(balance.total(), 1700); + assert_eq!(balance.total(), 1800); } #[test] #[should_panic(expected = "attempt to add with overflow")] fn test_balance_overflow() { - let balance = WalletBalance::new(u64::MAX, u64::MAX, u64::MAX); + let balance = WalletBalance::new(u64::MAX, u64::MAX, u64::MAX, u64::MAX); balance.total(); } #[test] fn test_balance_display() { let zero = WalletBalance::default(); - assert_eq!(zero.to_string(), "Spendable: 0, Unconfirmed: 0, Locked: 0, Total: 0"); + assert_eq!( + zero.to_string(), + "Spendable: 0, Unconfirmed: 0, Immature: 0, Locked: 0, Total: 0" + ); - let balance = WalletBalance::new(1000, 500, 200); + let balance = WalletBalance::new(1000, 500, 100, 200); let display = balance.to_string(); - assert_eq!(display, "Spendable: 1000, Unconfirmed: 500, Locked: 200, Total: 1700"); + assert_eq!( + display, + "Spendable: 1000, Unconfirmed: 500, Immature: 100, Locked: 200, Total: 1800" + ); } #[test] fn test_balance_add_assign() { - let mut balance = WalletBalance::new(1000, 500, 200); - let balance_add = WalletBalance::new(300, 100, 50); + let mut balance = WalletBalance::new(1000, 500, 50, 200); + let balance_add = WalletBalance::new(300, 100, 100, 50); // Test adding actual balances balance += balance_add; assert_eq!(balance.spendable(), 1300); assert_eq!(balance.unconfirmed(), 600); + assert_eq!(balance.immature(), 150); assert_eq!(balance.locked(), 250); - assert_eq!(balance.total(), 2150); + assert_eq!(balance.total(), 2300); // Test adding zero balances let balance_before = balance; balance += WalletBalance::default(); diff --git a/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs b/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs index 54c66f768..128e2a1a3 100644 --- a/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs +++ b/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs @@ -91,9 +91,6 @@ pub trait WalletInfoInterface: Sized + WalletTransactionChecker + ManagedAccount /// Get immature transactions fn immature_transactions(&self) -> Vec; - /// Get immature balance - fn immature_balance(&self) -> u64; - /// Create an unsigned payment transaction #[allow(clippy::too_many_arguments)] fn create_unsigned_payment_transaction( @@ -242,14 +239,6 @@ impl WalletInfoInterface for ManagedWalletInfo { transactions } - fn immature_balance(&self) -> u64 { - self.utxos() - .iter() - .filter(|utxo| utxo.is_coinbase && !utxo.is_mature(self.synced_height())) - .map(|utxo| utxo.value()) - .sum() - } - fn create_unsigned_payment_transaction( &mut self, wallet: &Wallet, @@ -272,5 +261,7 @@ impl WalletInfoInterface for ManagedWalletInfo { fn update_synced_height(&mut self, current_height: u32) { self.metadata.synced_height = current_height; + // Update cached balance + self.update_balance(); } }