Skip to content

Commit

Permalink
Change get_balance to return in categories.
Browse files Browse the repository at this point in the history
Add type balance with add, display traits. Change affected tests.
  • Loading branch information
wszdexdrf committed Jun 23, 2022
1 parent 93ca3b9 commit f04341b
Show file tree
Hide file tree
Showing 3 changed files with 105 additions and 11 deletions.
2 changes: 1 addition & 1 deletion src/blockchain/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down
47 changes: 47 additions & 0 deletions src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down
67 changes: 57 additions & 10 deletions src/wallet/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -4686,6 +4722,17 @@ 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;
Expand All @@ -4695,14 +4742,14 @@ pub(crate) mod test {
// 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(),
Expand All @@ -4715,7 +4762,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(),
Expand All @@ -4728,7 +4775,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();
}
Expand Down

0 comments on commit f04341b

Please sign in to comment.