feat: add balance method to Wallet

which in turn calls `Wallet::balance_with_params_conf_threshold` under
the hood. Also added some test utilities and a test to check the
implementation of the new methods.

Author: 110CodingP

src/test_utils.rs

@@ -2,6 +2,7 @@
 #![allow(unused)]
 use alloc::string::ToString;
 use alloc::sync::Arc;
+use core::fmt;
 use core::str::FromStr;
 
 use bdk_chain::{BlockId, ConfirmationBlockTime, TxUpdate};
@@ -303,50 +304,57 @@ impl From<ConfirmationBlockTime> for ReceiveTo {
 //     OutPoint { txid, vout: 0 }
 // }
 
-// /// Insert a checkpoint into the wallet. This can be used to extend the wallet's local chain
-// /// or to insert a block that did not exist previously. Note that if replacing a block with
-// /// a different one at the same height, then all later blocks are evicted as well.
-// pub fn insert_checkpoint(wallet: &mut Wallet, block: BlockId) {
-//     let mut cp = wallet.latest_checkpoint();
-//     cp = cp.insert(block);
-//     wallet
-//         .apply_update(Update {
-//             chain: Some(cp),
-//             ..Default::default()
-//         })
-//         .unwrap();
-// }
+/// Insert a checkpoint into the wallet. This can be used to extend the wallet's local chain
+/// or to insert a block that did not exist previously. Note that if replacing a block with
+/// a different one at the same height, then all later blocks are evicted as well.
+pub fn insert_checkpoint<K: Ord + Clone + fmt::Debug>(wallet: &mut Wallet<K>, block: BlockId) {
+    let mut cp = wallet.latest_checkpoint();
+    cp = cp.insert(block);
+    wallet
+        .apply_update(Update {
+            chain: Some(cp),
+            ..Default::default()
+        })
+        .unwrap();
+}
 
-// /// Inserts a transaction into the local view, assuming it is currently present in the mempool.
-// ///
-// /// This can be used, for example, to track a transaction immediately after it is broadcast.
-// pub fn insert_tx(wallet: &mut Wallet, tx: Transaction) {
-//     let txid = tx.compute_txid();
-//     let seen_at = std::time::UNIX_EPOCH.elapsed().unwrap().as_secs();
-//     let mut tx_update = TxUpdate::default();
-//     tx_update.txs = vec![Arc::new(tx)];
-//     tx_update.seen_ats = [(txid, seen_at)].into();
-//     wallet
-//         .apply_update(Update {
-//             tx_update,
-//             ..Default::default()
-//         })
-//         .expect("failed to apply update");
-// }
+/// Inserts a transaction into the local view, assuming it is currently present in the mempool.
+///
+/// This can be used, for example, to track a transaction immediately after it is broadcast.
+pub fn insert_tx<K>(wallet: &mut Wallet<K>, tx: Transaction)
+where
+    K: Ord + fmt::Debug + Clone,
+{
+    let txid = tx.compute_txid();
+    let seen_at = std::time::UNIX_EPOCH.elapsed().unwrap().as_secs();
+    let mut tx_update = TxUpdate::default();
+    tx_update.txs = vec![Arc::new(tx)];
+    tx_update.seen_ats = [(txid, seen_at)].into();
+    wallet
+        .apply_update(Update {
+            tx_update,
+            ..Default::default()
+        })
+        .expect("failed to apply update");
+}
 
-// /// Simulates confirming a tx with `txid` by applying an update to the wallet containing
-// /// the given `anchor`. Note: to be considered confirmed the anchor block must exist in
-// /// the current active chain.
-// pub fn insert_anchor(wallet: &mut Wallet, txid: Txid, anchor: ConfirmationBlockTime) {
-//     let mut tx_update = TxUpdate::default();
-//     tx_update.anchors = [(anchor, txid)].into();
-//     wallet
-//         .apply_update(Update {
-//             tx_update,
-//             ..Default::default()
-//         })
-//         .expect("failed to apply update");
-// }
+/// Simulates confirming a tx with `txid` by applying an update to the wallet containing
+/// the given `anchor`. Note: to be considered confirmed the anchor block must exist in
+/// the current active chain.
+pub fn insert_anchor<K: Ord + fmt::Debug + Clone>(
+    wallet: &mut Wallet<K>,
+    txid: Txid,
+    anchor: ConfirmationBlockTime,
+) {
+    let mut tx_update = TxUpdate::default();
+    tx_update.anchors = [(anchor, txid)].into();
+    wallet
+        .apply_update(Update {
+            tx_update,
+            ..Default::default()
+        })
+        .expect("failed to apply update");
+}
 
 // /// Marks the given `txid` seen as unconfirmed at `seen_at`
 // pub fn insert_seen_at(wallet: &mut Wallet, txid: Txid, seen_at: u64) {

src/wallet/mod.rs

@@ -34,7 +34,7 @@ use bdk_chain::{
         SyncResponse,
     },
     tx_graph::{self, CalculateFeeError, CanonicalTx, TxGraph, TxUpdate},
-    BlockId, CanonicalizationParams, ChainPosition, ConfirmationBlockTime, DescriptorExt,
+    Anchor, BlockId, CanonicalizationParams, ChainPosition, ConfirmationBlockTime, DescriptorExt,
     FullTxOut, Indexed, IndexedTxGraph, Indexer, Merge,
 };
 use bitcoin::{
@@ -571,6 +571,107 @@ where
     }
 }
 
+impl<K> Wallet<K>
+where
+    K: Ord + Clone + Debug,
+{
+    /// Computes the wallet balance.
+    pub fn balance_with_params_conf_threshold(
+        &self,
+        params: CanonicalizationParams,
+        outpoints: impl IntoIterator<Item = ((K, u32), OutPoint)>,
+        conf_threshold: u32,
+        trust_predicate: impl Fn(&FullTxOut<ConfirmationBlockTime>) -> bool,
+    ) -> Balance {
+        let mut immature = Amount::ZERO;
+        let mut trusted_pending = Amount::ZERO;
+        let mut untrusted_pending = Amount::ZERO;
+        let mut confirmed = Amount::ZERO;
+
+        let chain = &self.chain;
+        let chain_tip = chain.tip().block_id();
+
+        for (_, txo) in self
+            .tx_graph
+            .graph()
+            .filter_chain_unspents(chain, chain_tip, params, outpoints)
+        {
+            match &txo.chain_position {
+                ChainPosition::Confirmed { anchor, .. } => {
+                    let confirmation_height = anchor.confirmation_height_upper_bound();
+                    let confirmations = chain_tip
+                        .height
+                        .saturating_sub(confirmation_height)
+                        .saturating_add(1);
+
+                    if confirmations < conf_threshold {
+                        if trust_predicate(&txo) {
+                            trusted_pending += txo.txout.value;
+                        } else {
+                            untrusted_pending += txo.txout.value;
+                        }
+                    } else if txo.is_confirmed_and_spendable(chain_tip.height) {
+                        confirmed += txo.txout.value;
+                    } else if !txo.is_mature(chain_tip.height) {
+                        immature += txo.txout.value;
+                    }
+                }
+                ChainPosition::Unconfirmed { .. } => {
+                    if trust_predicate(&txo) {
+                        trusted_pending += txo.txout.value;
+                    } else {
+                        untrusted_pending += txo.txout.value;
+                    }
+                }
+            }
+        }
+
+        Balance {
+            immature,
+            trusted_pending,
+            untrusted_pending,
+            confirmed,
+        }
+    }
+
+    /// Computes the wallet balance.
+    pub fn balance(&self) -> Balance {
+        self.balance_with_params_conf_threshold(
+            CanonicalizationParams::default(),
+            self.tx_graph.index.outpoints().clone(),
+            1,
+            |txo| self.is_tx_trusted(txo.outpoint.txid),
+        )
+    }
+
+    /// Computes the wallet balance over a keychain range.
+    pub fn keychain_balance(&self, keychains: impl core::ops::RangeBounds<K>) -> Balance {
+        self.balance_with_params_conf_threshold(
+            CanonicalizationParams::default(),
+            self.tx_graph.index.keychain_outpoints_in_range(keychains),
+            1,
+            |txo| self.is_tx_trusted(txo.outpoint.txid),
+        )
+    }
+
+    /// Whether the transaction of `txid` is considered trusted by this wallet.
+    ///
+    /// Trust is defined as a tx of which all of the inputs are controlled by the wallet, as we can
+    /// assume it won't be double-spent unintentionally.
+    pub fn is_tx_trusted(&self, txid: Txid) -> bool {
+        let Some(tx) = self.tx_graph.graph().get_tx(txid) else {
+            return false;
+        };
+        if tx.input.is_empty() {
+            return false;
+        }
+        tx.input.iter().all(|txin| {
+            let outpoint = txin.previous_output;
+            self.tx_graph.index.txout(outpoint).is_some()
+        })
+    }
+}
+
 // TODO: replace with `PersistedWallet`
 #[cfg(feature = "rusqlite")]
 impl<K> Wallet<K>