Change get_balance to return in categories.

Add type balance with add, display traits. Change affected tests.

Author: wszdexdrf

src/blockchain/mod.rs

@@ -187,7 +187,7 @@ This example shows how to sync multiple walles and return the sum of their balan
 # use bdk::database::*;
 # use bdk::wallet::*;
 # use bdk::*;
-fn sum_of_balances<B: BlockchainFactory>(blockchain_factory: B, wallets: &[Wallet<MemoryDatabase>]) -> Result<u64, Error> {
+fn sum_of_balances<B: BlockchainFactory>(blockchain_factory: B, wallets: &[Wallet<MemoryDatabase>]) -> Result<Balance, Error> {
     Ok(wallets
         .iter()
         .map(|w| -> Result<_, Error> {

src/types.rs

@@ -240,6 +240,53 @@ impl BlockTime {
     }
 }
 
+/// Balance differentiated in various categories
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
+pub struct Balance {
+    /// All coinbase not matured yet
+    pub immature: u64,
+    /// In wallet change utxo, but not confirmed yet.
+    pub trusted_pending: u64,
+    /// Other wallet received utxo not confirmed yet.
+    pub untrusted_pending: u64,
+    /// All rest available balance
+    pub available: u64,
+}
+
+impl std::fmt::Display for Balance {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{{ immature: {}, trusted_pending: {}, untrusted_pending: {}, available: {} }}",
+            self.immature, self.trusted_pending, self.untrusted_pending, self.available
+        )
+    }
+}
+
+impl std::ops::Add for Balance {
+    type Output = Self;
+
+    fn add(self, other: Self) -> Self {
+        Self {
+            immature: self.immature + other.immature,
+            trusted_pending: self.trusted_pending + other.trusted_pending,
+            untrusted_pending: self.untrusted_pending + other.untrusted_pending,
+            available: self.available + other.available,
+        }
+    }
+}
+
+impl std::iter::Sum for Balance {
+    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
+        iter.fold(
+            Balance {
+                ..Default::default()
+            },
+            |a, b| a + b,
+        )
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/wallet/mod.rs

@@ -446,15 +446,51 @@ where
         self.database.borrow().iter_txs(include_raw)
     }
 
-    /// Return the balance, meaning the sum of this wallet's unspent outputs' values
+    /// Return the balance, separated into available, trusted-pending, untrusted-pending and immature
+    /// values.
     ///
     /// Note that this methods only operate on the internal database, which first needs to be
     /// [`Wallet::sync`] manually.
-    pub fn get_balance(&self) -> Result<u64, Error> {
-        Ok(self
-            .list_unspent()?
-            .iter()
-            .fold(0, |sum, i| sum + i.txout.value))
+    pub fn get_balance(&self) -> Result<Balance, Error> {
+        let mut immature = 0;
+        let mut trusted_pending = 0;
+        let mut untrusted_pending = 0;
+        let mut available = 0;
+        let utxos = self.list_unspent()?;
+        let database = self.database.borrow();
+        for u in utxos {
+            // Unwrap used since utxo set is created from database
+            let tx = database.get_tx(&u.outpoint.txid, true)?.unwrap();
+            if let Some(tx_conf_time) = &tx.confirmation_time {
+                if tx.transaction.expect("No transaction").is_coin_base() {
+                    let last_sync_height = self
+                        .database()
+                        .get_sync_time()?
+                        .map(|sync_time| sync_time.block_time.height);
+
+                    if let Some(last_sync_height) = last_sync_height {
+                        if (last_sync_height - tx_conf_time.height) >= COINBASE_MATURITY {
+                            available += u.txout.value;
+                        } else {
+                            immature += u.txout.value;
+                        }
+                    }
+                } else {
+                    available += u.txout.value;
+                }
+            } else if u.keychain == KeychainKind::Internal {
+                trusted_pending += u.txout.value;
+            } else {
+                untrusted_pending += u.txout.value;
+            }
+        }
+
+        Ok(Balance {
+            immature,
+            trusted_pending,
+            untrusted_pending,
+            available,
+        })
     }
 
     /// Add an external signer
@@ -4686,23 +4722,30 @@ pub(crate) mod test {
             Some(confirmation_time),
             (@coinbase true)
         );
+        let sync_time = SyncTime {
+            block_time: BlockTime {
+                height: confirmation_time,
+                timestamp: 0,
+            },
+        };
+        wallet
+            .database
+            .borrow_mut()
+            .set_sync_time(sync_time.clone())
+            .unwrap();
 
         let not_yet_mature_time = confirmation_time + COINBASE_MATURITY - 1;
         let maturity_time = confirmation_time + COINBASE_MATURITY;
 
-        // The balance is nonzero, even if we can't spend anything
-        // FIXME: we should differentiate the balance between immature,
-        // trusted, untrusted_pending
-        // See https://github.com/bitcoindevkit/bdk/issues/238
         let balance = wallet.get_balance().unwrap();
-        assert!(balance != 0);
+        assert!(balance.immature != 0);
 
         // We try to create a transaction, only to notice that all
         // our funds are unspendable
         let addr = Address::from_str("2N1Ffz3WaNzbeLFBb51xyFMHYSEUXcbiSoX").unwrap();
         let mut builder = wallet.build_tx();
         builder
-            .add_recipient(addr.script_pubkey(), balance / 2)
+            .add_recipient(addr.script_pubkey(), balance.immature / 2)
             .set_current_height(confirmation_time);
         assert!(matches!(
             builder.finish().unwrap_err(),
@@ -4715,7 +4758,7 @@ pub(crate) mod test {
         // Still unspendable...
         let mut builder = wallet.build_tx();
         builder
-            .add_recipient(addr.script_pubkey(), balance / 2)
+            .add_recipient(addr.script_pubkey(), balance.immature / 2)
             .set_current_height(not_yet_mature_time);
         assert!(matches!(
             builder.finish().unwrap_err(),
@@ -4728,7 +4771,7 @@ pub(crate) mod test {
         // ...Now the coinbase is mature :)
         let mut builder = wallet.build_tx();
         builder
-            .add_recipient(addr.script_pubkey(), balance / 2)
+            .add_recipient(addr.script_pubkey(), balance.immature / 2)
             .set_current_height(maturity_time);
         builder.finish().unwrap();
     }