Skip to content

Commit aa61e72

Browse files
committed
refactor: cleanup WalletBalance
- some renaming of the fields - make the fields private - drop overflow checks and restrict total to `u64::MAX` - implement, test and use `core::fmt::Display` - simplify tests - drop unused/test-only code
1 parent cd12793 commit aa61e72

File tree

17 files changed

+78
-404
lines changed

17 files changed

+78
-404
lines changed

dash-spv/src/main.rs

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -451,7 +451,7 @@ async fn run_client<S: dash_spv::storage::StorageManager>(
451451
_ = tokio::time::sleep(snapshot_interval) => {
452452
// Log snapshot if interval has elapsed
453453
if last_snapshot.elapsed() >= snapshot_interval {
454-
let (tx_count, confirmed, unconfirmed, locked, total) = {
454+
let (tx_count, wallet_balance) = {
455455
let mgr = wallet_for_logger.read().await;
456456

457457
// Count wallet-affecting transactions from wallet transaction history
@@ -461,20 +461,16 @@ async fn run_client<S: dash_spv::storage::StorageManager>(
461461
.unwrap_or(0);
462462

463463
// Read wallet balance from the managed wallet info
464-
let wb = mgr.get_wallet_balance(&wallet_id_for_logger).ok();
465-
let (c, u, l, t) = wb.map(|b| (b.confirmed, b.unconfirmed, b.locked, b.total)).unwrap_or((0, 0, 0, 0));
464+
let wallet_balance = mgr.get_wallet_balance(&wallet_id_for_logger).unwrap_or_default();
466465

467-
(tx_count, c, u, l, t)
466+
(tx_count, wallet_balance)
468467
};
469468
tracing::info!(
470-
"Wallet tx summary: tx_count={} (blocks={} + mempool={}), balances: confirmed={} unconfirmed={} locked={} total={}",
469+
"Wallet tx summary: tx_count={} (blocks={} + mempool={}), balances: {}",
471470
tx_count,
472471
total_detected_block_txs,
473472
total_detected_mempool_txs,
474-
confirmed,
475-
unconfirmed,
476-
locked,
477-
total,
473+
wallet_balance,
478474
);
479475
last_snapshot = std::time::Instant::now();
480476
}

key-wallet-ffi/src/managed_account.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -513,10 +513,10 @@ pub unsafe extern "C" fn managed_account_get_balance(
513513
let balance = &account.inner().balance;
514514

515515
*balance_out = crate::types::FFIBalance {
516-
confirmed: balance.confirmed,
517-
unconfirmed: balance.unconfirmed,
516+
confirmed: balance.spendable(),
517+
unconfirmed: balance.unconfirmed(),
518518
immature: 0, // WalletBalance doesn't have immature field
519-
total: balance.total,
519+
total: balance.total(),
520520
};
521521

522522
true

key-wallet-ffi/src/managed_wallet.rs

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -576,10 +576,10 @@ pub unsafe extern "C" fn managed_wallet_get_balance(
576576
let balance = &managed_wallet.inner().balance;
577577

578578
unsafe {
579-
*confirmed_out = balance.confirmed;
580-
*unconfirmed_out = balance.unconfirmed;
581-
*locked_out = balance.locked;
582-
*total_out = balance.total;
579+
*confirmed_out = balance.spendable();
580+
*unconfirmed_out = balance.unconfirmed();
581+
*locked_out = balance.locked();
582+
*total_out = balance.total();
583583
}
584584

585585
FFIError::set_success(error);
@@ -1042,12 +1042,7 @@ mod tests {
10421042
let mut managed_info = ManagedWalletInfo::from_wallet(wallet_arc);
10431043

10441044
// Set some test balance values
1045-
managed_info.balance = WalletBalance {
1046-
confirmed: 1000000,
1047-
unconfirmed: 50000,
1048-
locked: 25000,
1049-
total: 1075000,
1050-
};
1045+
managed_info.balance = WalletBalance::new(1000000, 50000, 25000);
10511046

10521047
let ffi_managed = FFIManagedWalletInfo::new(managed_info);
10531048
let ffi_managed_ptr = Box::into_raw(Box::new(ffi_managed));

key-wallet-ffi/src/types.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,10 +64,10 @@ pub struct FFIBalance {
6464
impl From<key_wallet::WalletBalance> for FFIBalance {
6565
fn from(balance: key_wallet::WalletBalance) -> Self {
6666
FFIBalance {
67-
confirmed: balance.confirmed,
68-
unconfirmed: balance.unconfirmed,
69-
immature: balance.locked, // Map locked to immature for now
70-
total: balance.total,
67+
confirmed: balance.spendable(),
68+
unconfirmed: balance.unconfirmed(),
69+
immature: balance.locked(), // Map locked to immature for now
70+
total: balance.total(),
7171
}
7272
}
7373
}

key-wallet-ffi/src/wallet_manager.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -703,8 +703,8 @@ pub unsafe extern "C" fn wallet_manager_get_wallet_balance(
703703
match result {
704704
Ok(balance) => {
705705
unsafe {
706-
*confirmed_out = balance.confirmed;
707-
*unconfirmed_out = balance.unconfirmed;
706+
*confirmed_out = balance.spendable();
707+
*unconfirmed_out = balance.unconfirmed();
708708
}
709709
FFIError::set_success(error);
710710
true

key-wallet-manager/examples/wallet_creation.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ fn main() {
127127
for (i, wallet_id) in [wallet_id, wallet_id2].iter().enumerate() {
128128
match manager.get_wallet_balance(wallet_id) {
129129
Ok(balance) => {
130-
println!(" Wallet {}: {} satoshis", i + 1, balance.total);
130+
println!(" Wallet {}: {} satoshis", i + 1, balance.total());
131131
}
132132
Err(e) => {
133133
println!(" Wallet {}: Error - {:?}", i + 1, e);

key-wallet-manager/src/wallet_manager/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -880,7 +880,7 @@ impl<T: WalletInfoInterface> WalletManager<T> {
880880

881881
/// Get total balance across all wallets and networks
882882
pub fn get_total_balance(&self) -> u64 {
883-
self.wallet_infos.values().map(|info| info.balance().total).sum()
883+
self.wallet_infos.values().map(|info| info.balance().total()).sum()
884884
}
885885

886886
/// Get balance for a specific wallet

key-wallet-manager/tests/integration_test.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ fn test_utxo_management() {
126126

127127
let balance = manager.get_wallet_balance(&wallet_id);
128128
assert!(balance.is_ok());
129-
assert_eq!(balance.unwrap().total, 0);
129+
assert_eq!(balance.unwrap().total(), 0);
130130
}
131131

132132
#[test]
@@ -145,7 +145,7 @@ fn test_balance_calculation() {
145145
// Check wallet balance (should be 0 initially)
146146
let balance = manager.get_wallet_balance(&wallet_id);
147147
assert!(balance.is_ok());
148-
assert_eq!(balance.unwrap().total, 0);
148+
assert_eq!(balance.unwrap().total(), 0);
149149

150150
// Check global balance
151151
let total = manager.get_total_balance();

key-wallet/src/lib.rs

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,7 @@ pub use managed_account::managed_account_type::ManagedAccountType;
6565
pub use mnemonic::Mnemonic;
6666
pub use seed::Seed;
6767
pub use utxo::{Utxo, UtxoSet};
68-
pub use wallet::{
69-
balance::{BalanceError, WalletBalance},
70-
Wallet,
71-
};
68+
pub use wallet::{balance::WalletBalance, Wallet};
7269

7370
/// Re-export commonly used types
7471
pub mod prelude {

key-wallet/src/managed_account/mod.rs

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -265,15 +265,9 @@ impl ManagedAccount {
265265
}
266266

267267
/// Update the account balance
268-
pub fn update_balance(
269-
&mut self,
270-
confirmed: u64,
271-
unconfirmed: u64,
272-
locked: u64,
273-
) -> Result<(), crate::wallet::balance::BalanceError> {
274-
self.balance.update(confirmed, unconfirmed, locked)?;
268+
pub fn update_balance(&mut self, confirmed: u64, unconfirmed: u64, locked: u64) {
269+
self.balance = WalletBalance::new(confirmed, unconfirmed, locked);
275270
self.metadata.last_used = Some(Self::current_timestamp());
276-
Ok(())
277271
}
278272

279273
/// Get all addresses from all pools

0 commit comments

Comments
 (0)