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

@@ -571,6 +571,157 @@ where
     }
 }
 
+impl<K> Wallet<K>
+where
+    K: Ord + Clone + Debug,
+{
+    /// Compute the wallet balance with canonical params, confirmation threshold, and trust
+    /// predicate.
+    ///
+    /// Panics if `conf_threshold` is equal to 0.
+    pub fn balance_with_params_conf_threshold(
+        &self,
+        params: CanonicalizationParams,
+        conf_threshold: u32,
+        trust_predicate: impl Fn(
+            &OutPoint,
+            &HashMap<Txid, CanonicalTx<Arc<Transaction>, ConfirmationBlockTime>>,
+        ) -> bool,
+    ) -> Balance {
+        use crate::chain::ChainOracle;
+        let mut confirmed = Amount::ZERO;
+        let mut trusted_pending = Amount::ZERO;
+        let mut untrusted_pending = Amount::ZERO;
+        let mut immature = Amount::ZERO;
+
+        let mut canon_txs = HashMap::new();
+        let mut canon_spends = HashMap::new();
+        let outpoints = self.tx_graph.index.outpoints().iter().cloned();
+
+        for res in self.tx_graph.graph().try_list_canonical_txs(
+            &self.chain,
+            self.chain.tip().block_id(),
+            params,
+        ) {
+            let canonical_tx = res.expect("oracle is infallible");
+            let txid = canonical_tx.tx_node.txid;
+
+            if !canonical_tx.tx_node.is_coinbase() {
+                for txin in &canonical_tx.tx_node.tx.input {
+                    let _res = canon_spends.insert(txin.previous_output, txid);
+                    assert!(_res.is_none(), "tried to replace {_res:?} with {txid:?} ")
+                }
+            }
+
+            canon_txs.insert(txid, canonical_tx);
+        }
+
+        let unspent_txouts = outpoints.into_iter().filter_map(|(_, outpoint)| {
+            if canon_spends.contains_key(&outpoint) {
+                return None;
+            }
+            let canon_tx = canon_txs.get(&outpoint.txid)?;
+            let txout = canon_tx
+                .tx_node
+                .tx
+                .output
+                .get(outpoint.vout as usize)
+                .cloned()
+                .expect("oracle is infallible");
+            let chain_position = canon_tx.chain_position;
+            let is_on_coinbase = canon_tx.tx_node.is_coinbase();
+            Some(FullTxOut {
+                chain_position,
+                outpoint,
+                is_on_coinbase,
+                spent_by: None,
+                txout,
+            })
+        });
+
+        let target_height = self.chain.tip().height().checked_sub(
+            conf_threshold
+                .checked_sub(1)
+                .expect("conf threshold should be positive integer"),
+        );
+        let curr_height = self.chain.tip().height();
+
+        for full_txout in unspent_txouts {
+            match full_txout.chain_position {
+                ChainPosition::Confirmed { .. } => match target_height {
+                    Some(ht) => {
+                        if full_txout.is_confirmed_and_spendable(ht) {
+                            confirmed += full_txout.txout.value;
+                        } else if full_txout.is_confirmed_and_spendable(curr_height) {
+                            if full_txout.is_on_coinbase {
+                                confirmed += full_txout.txout.value;
+                            } else if trust_predicate(&full_txout.outpoint, &canon_txs) {
+                                trusted_pending += full_txout.txout.value;
+                            } else {
+                                untrusted_pending += full_txout.txout.value;
+                            }
+                        } else if !full_txout.is_mature(curr_height) {
+                            immature += full_txout.txout.value;
+                        }
+                    }
+                    None => {
+                        if full_txout.is_confirmed_and_spendable(curr_height) {
+                            if full_txout.is_on_coinbase {
+                                confirmed += full_txout.txout.value;
+                            } else if trust_predicate(&full_txout.outpoint, &canon_txs) {
+                                trusted_pending += full_txout.txout.value;
+                            } else {
+                                untrusted_pending += full_txout.txout.value;
+                            }
+                        } else if !full_txout.is_mature(curr_height) {
+                            immature += full_txout.txout.value;
+                        }
+                    }
+                },
+                ChainPosition::Unconfirmed { .. } => {
+                    if trust_predicate(&full_txout.outpoint, &canon_txs) {
+                        trusted_pending += full_txout.txout.value;
+                    } else {
+                        untrusted_pending += full_txout.txout.value;
+                    }
+                }
+            }
+        }
+
+        Balance {
+            confirmed,
+            trusted_pending,
+            untrusted_pending,
+            immature,
+        }
+    }
+
+    /// Compute the wallet balance with the default parameters.
+    pub fn balance(&self) -> Balance {
+        self.balance_with_params_conf_threshold(
+            CanonicalizationParams::default(),
+            1,
+            |outpoint, canon_txs| {
+                let mut trusted = true;
+                let canon_tx = canon_txs.get(&outpoint.txid).expect("oracle is infallible");
+                for txin in &canon_tx.tx_node.tx.input {
+                    trusted = trusted
+                        && self
+                            .tx_graph
+                            .index
+                            .outpoints()
+                            .iter()
+                            .any(|(_, item)| *item == txin.previous_output);
+                }
+                if canon_tx.tx_node.tx.input.is_empty() {
+                    trusted = false;
+                }
+                trusted
+            },
+        )
+    }
+}
+
 // TODO: replace with `PersistedWallet`
 #[cfg(feature = "rusqlite")]
 impl<K> Wallet<K>