fix: unify immature balance tracking with remaining balances

We currently keep track of immature balance separately from other balances with the `immature_balance` function. This PR drops the extra function, adds `WalletBalance.immature` and changes the balance tracking so that we have all balances in one place. It also fixes a mix-up between locked and immature balance in FFI.

Author: xdustinface

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`
 

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)

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

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();
     }
@@ -1041,15 +1045,16 @@ mod tests {
         let wallet_arc = unsafe { &(*wallet_ptr).wallet };
         let mut managed_info = ManagedWalletInfo::from_wallet(wallet_arc);
 
-        // Set some test balance values
-        managed_info.balance = WalletBalance::new(1000000, 50000, 25000);
+        // Set some test balance values (confirmed, unconfirmed, immature, locked)
+        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));
 
         // 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,15 +1073,17 @@ 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 {
             managed_wallet_get_balance(
                 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,

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<key_wallet::WalletBalance> 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(),
         }
     }

key-wallet-manager/tests/spv_integration_tests.rs

@@ -162,6 +162,29 @@ fn assert_wallet_heights(manager: &WalletManager<ManagedWalletInfo>, 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() {
@@ -181,3 +204,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::<ManagedWalletInfo>::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 + 100;
+    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"
+    );
+}

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());
     }
 

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)
+    }
+}

key-wallet/src/test_utils/mod.rs

@@ -1,3 +1,4 @@
+mod account;
 mod address;
 mod utxo;