Merge #640: Get balance in categories

0f03831 Change get_balance to return in categories. (wszdexdrf)

Pull request description:

  ### Description
  This changes `get_balance()` function so that it returns balance separated in 4 categories:
  - available
  - trusted-pending
  - untrusted-pending
  - immature

  Fixes #238

  ### Notes to the reviewers
  Based on #614

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo fmt` and `cargo clippy` before committing

  #### New Features:

  * [x] I've updated tests for the new feature
  * [x] I've added docs for the new feature
  * [x] I've updated `CHANGELOG.md`

ACKs for top commit:
  afilini:
    ACK 0f03831

Tree-SHA512: 39f02c22c61b6c73dd8e6d27b1775a72e64ab773ee67c0ad00e817e555c52cdf648f482ca8be5fcc2f3d62134c35b720b1e61b311cb6debb3ad651e79c829b93

Author: afilini

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Consolidate params `fee_amount` and `amount_needed` in `target_amount` in `CoinSelectionAlgorithm::coin_select` signature.
 - Change the meaning of the `fee_amount` field inside `CoinSelectionResult`: from now on the `fee_amount` will represent only the fees asociated with the utxos in the `selected` field of `CoinSelectionResult`.
 - New `RpcBlockchain` implementation with various fixes.
+- Return balance in separate categories, namely `confirmed`, `trusted_pending`, `untrusted_pending` & `immature`.
 
 ## [v0.20.0] - [v0.19.0]
 

src/blockchain/electrum.rs

@@ -385,7 +385,7 @@ mod test {
             .sync_wallet(&wallet, None, Default::default())
             .unwrap();
 
-        assert_eq!(wallet.get_balance().unwrap(), 50_000);
+        assert_eq!(wallet.get_balance().unwrap().untrusted_pending, 50_000);
     }
 
     #[test]

src/blockchain/mod.rs

@@ -194,7 +194,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/testutils/blockchain_tests.rs

@@ -124,7 +124,7 @@ where
         // perform wallet sync
         wallet.sync(&blockchain, Default::default()).unwrap();
 
-        let wallet_balance = wallet.get_balance().unwrap();
+        let wallet_balance = wallet.get_balance().unwrap().get_total();
         println!(
             "max: {}, min: {}, actual: {}",
             max_balance, min_balance, wallet_balance
@@ -193,7 +193,7 @@ where
     wallet.sync(&blockchain, Default::default()).unwrap();
     println!("sync done!");
 
-    let balance = wallet.get_balance().unwrap();
+    let balance = wallet.get_balance().unwrap().get_total();
     assert_eq!(balance, expected_balance);
 }
 
@@ -245,13 +245,13 @@ where
 
     // actually test the wallet
     wallet.sync(&blockchain, Default::default()).unwrap();
-    let balance = wallet.get_balance().unwrap();
+    let balance = wallet.get_balance().unwrap().get_total();
     assert_eq!(balance, expected_balance);
 
     // now try with a fresh wallet
     let fresh_wallet =
         Wallet::new(descriptor, None, Network::Regtest, MemoryDatabase::new()).unwrap();
     fresh_wallet.sync(&blockchain, Default::default()).unwrap();
-    let fresh_balance = fresh_wallet.get_balance().unwrap();
+    let fresh_balance = fresh_wallet.get_balance().unwrap().get_total();
     assert_eq!(fresh_balance, expected_balance);
 }

src/testutils/configurable_blockchain_tests.rs

@@ -227,7 +227,7 @@ pub struct TransactionDetails {
     /// Sent value (sats)
     /// Sum of owned inputs of this transaction.
     pub sent: u64,
-    /// Fee value (sats) if available.
+    /// Fee value (sats) if confirmed.
     /// The availability of the fee depends on the backend. It's never `None` with an Electrum
     /// Server backend, but it could be `None` with a Bitcoin RPC node without txindex that receive
     /// funds while offline.
@@ -262,6 +262,65 @@ impl BlockTime {
     }
 }
 
+/// Balance differentiated in various categories
+#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Default)]
+pub struct Balance {
+    /// All coinbase outputs not yet matured
+    pub immature: u64,
+    /// Unconfirmed UTXOs generated by a wallet tx
+    pub trusted_pending: u64,
+    /// Unconfirmed UTXOs received from an external wallet
+    pub untrusted_pending: u64,
+    /// Confirmed and immediately spendable balance
+    pub confirmed: u64,
+}
+
+impl Balance {
+    /// Get sum of trusted_pending and confirmed coins
+    pub fn get_spendable(&self) -> u64 {
+        self.confirmed + self.trusted_pending
+    }
+
+    /// Get the whole balance visible to the wallet
+    pub fn get_total(&self) -> u64 {
+        self.confirmed + self.trusted_pending + self.untrusted_pending + self.immature
+    }
+}
+
+impl std::fmt::Display for Balance {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{{ immature: {}, trusted_pending: {}, untrusted_pending: {}, confirmed: {} }}",
+            self.immature, self.trusted_pending, self.untrusted_pending, self.confirmed
+        )
+    }
+}
+
+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,
+            confirmed: self.confirmed + other.confirmed,
+        }
+    }
+}
+
+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/types.rs

@@ -460,15 +460,52 @@ 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 confirmed = 0;
+        let utxos = self.list_unspent()?;
+        let database = self.database.borrow();
+        let last_sync_height = match database
+            .get_sync_time()?
+            .map(|sync_time| sync_time.block_time.height)
+        {
+            Some(height) => height,
+            // None means database was never synced
+            None => return Ok(Balance::default()),
+        };
+        for u in utxos {
+            // Unwrap used since utxo set is created from database
+            let tx = database
+                .get_tx(&u.outpoint.txid, true)?
+                .expect("Transaction not found in database");
+            if let Some(tx_conf_time) = &tx.confirmation_time {
+                if tx.transaction.expect("No transaction").is_coin_base()
+                    && (last_sync_height - tx_conf_time.height) < COINBASE_MATURITY
+                {
+                    immature += u.txout.value;
+                } else {
+                    confirmed += 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,
+            confirmed,
+        })
     }
 
     /// Add an external signer
@@ -5232,23 +5269,38 @@ 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)
+            .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_eq!(
+            balance,
+            Balance {
+                immature: 25_000,
+                trusted_pending: 0,
+                untrusted_pending: 0,
+                confirmed: 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)
             .current_height(confirmation_time);
         assert!(matches!(
             builder.finish().unwrap_err(),
@@ -5261,7 +5313,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)
             .current_height(not_yet_mature_time);
         assert!(matches!(
             builder.finish().unwrap_err(),
@@ -5272,9 +5324,31 @@ pub(crate) mod test {
         ));
 
         // ...Now the coinbase is mature :)
+        let sync_time = SyncTime {
+            block_time: BlockTime {
+                height: maturity_time,
+                timestamp: 0,
+            },
+        };
+        wallet
+            .database
+            .borrow_mut()
+            .set_sync_time(sync_time)
+            .unwrap();
+
+        let balance = wallet.get_balance().unwrap();
+        assert_eq!(
+            balance,
+            Balance {
+                immature: 0,
+                trusted_pending: 0,
+                untrusted_pending: 0,
+                confirmed: 25_000
+            }
+        );
         let mut builder = wallet.build_tx();
         builder
-            .add_recipient(addr.script_pubkey(), balance / 2)
+            .add_recipient(addr.script_pubkey(), balance.confirmed / 2)
             .current_height(maturity_time);
         builder.finish().unwrap();
     }