From 7e8f849d4708efa6102d5c25a7647c2b47b90339 Mon Sep 17 00:00:00 2001
From: LLFourn <lloyd.fourn@gmail.com>
Date: Wed, 15 Feb 2023 13:00:32 +1100
Subject: [PATCH] Use bdk_chain Balance

instead of bdk's one
---
 src/types.rs      | 59 -----------------------------------------------
 src/wallet/mod.rs | 45 ++++--------------------------------
 tests/wallet.rs   |  4 +---
 3 files changed, 6 insertions(+), 102 deletions(-)

diff --git a/src/types.rs b/src/types.rs
index db33ed224..4472508f7 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -263,65 +263,6 @@ impl Ord for TransactionDetails {
     }
 }
 
-/// 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 core::fmt::Display for Balance {
-    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        write!(
-            f,
-            "{{ immature: {}, trusted_pending: {}, untrusted_pending: {}, confirmed: {} }}",
-            self.immature, self.trusted_pending, self.untrusted_pending, self.confirmed
-        )
-    }
-}
-
-impl core::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 core::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::*;
diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
index fe0660b6e..0f00dba98 100644
--- a/src/wallet/mod.rs
+++ b/src/wallet/mod.rs
@@ -19,6 +19,7 @@ use alloc::{
     sync::Arc,
     vec::Vec,
 };
+pub use bdk_chain::keychain::Balance;
 use bdk_chain::{
     chain_graph,
     keychain::{KeychainChangeSet, KeychainScan, KeychainTracker},
@@ -542,46 +543,10 @@ impl<D> Wallet<D> {
     /// Note that this method only operates on the internal database, which first needs to be
     /// [`Wallet::sync`] manually.
     pub fn get_balance(&self) -> Balance {
-        let mut immature = 0;
-        let mut trusted_pending = 0;
-        let mut untrusted_pending = 0;
-        let mut confirmed = 0;
-        let last_sync_height = match self.keychain_tracker.chain().latest_checkpoint() {
-            Some(last_sync_height) => last_sync_height.height,
-            // TODO: update this when we're allowed to add arbitary stuff to the mempool
-            None => return Balance::default(),
-        };
-
-        for ((keychain, _), utxo) in self.keychain_tracker.full_utxos() {
-            let confirmation_time = utxo.chain_position;
-            let is_coinbase = self
-                .keychain_tracker
-                .graph()
-                .get_tx(utxo.outpoint.txid)
-                .expect("must exist")
-                .is_coin_base();
-
-            match confirmation_time {
-                ConfirmationTime::Confirmed { height, .. } => {
-                    if is_coinbase && last_sync_height - height < COINBASE_MATURITY {
-                        immature += utxo.txout.value;
-                    } else {
-                        confirmed += utxo.txout.value;
-                    }
-                }
-                ConfirmationTime::Unconfirmed => match keychain {
-                    KeychainKind::External => untrusted_pending += utxo.txout.value,
-                    KeychainKind::Internal => trusted_pending += utxo.txout.value,
-                },
-            }
-        }
-
-        Balance {
-            immature,
-            trusted_pending,
-            untrusted_pending,
-            confirmed,
-        }
+        self.keychain_tracker.balance(|keychain| match keychain {
+            KeychainKind::External => false,
+            KeychainKind::Internal => true,
+        })
     }
 
     /// Add an external signer
diff --git a/tests/wallet.rs b/tests/wallet.rs
index 9e2d06852..44d489f9b 100644
--- a/tests/wallet.rs
+++ b/tests/wallet.rs
@@ -3,9 +3,7 @@ use bdk::descriptor::calc_checksum;
 use bdk::signer::{SignOptions, SignerError};
 use bdk::wallet::coin_selection::LargestFirstCoinSelection;
 use bdk::wallet::AddressIndex::*;
-use bdk::wallet::AddressInfo;
-use bdk::wallet::{AddressIndex, Wallet};
-use bdk::Balance;
+use bdk::wallet::{AddressIndex, AddressInfo, Balance, Wallet};
 use bdk::Error;
 use bdk::FeeRate;
 use bdk::KeychainKind;