diff --git a/src/electrum.rs b/src/electrum.rs index dd9af35..3da7bc4 100644 --- a/src/electrum.rs +++ b/src/electrum.rs @@ -226,19 +226,16 @@ impl Rpc { (scripthash,): (ScriptHash,), ) -> Result { let balance = match client.scripthashes.get(&scripthash) { - Some(status) => self.tracker.get_balance(status, &self.cache), + Some(status) => self.tracker.get_balance(&status), None => { warn!( "blockchain.scripthash.get_balance called for unsubscribed scripthash: {}", scripthash ); - self.tracker - .get_balance(&self.new_status(scripthash)?, &self.cache) + self.tracker.get_balance(&self.new_status(scripthash)?) } }; - Ok( - json!({"confirmed": balance.confirmed.as_sat(), "unconfirmed": balance.mempool_delta.as_sat()}), - ) + Ok(json!(balance)) } fn scripthash_get_history( diff --git a/src/status.rs b/src/status.rs index 1d48482..ba1edf2 100644 --- a/src/status.rs +++ b/src/status.rs @@ -41,14 +41,9 @@ impl TxEntry { } /// Relevant (scripthash-wise) funded outpoints - fn funding(&self) -> impl Iterator + '_ { + fn funding_outpoints(&self) -> impl Iterator + '_ { make_outpoints(&self.txid, &self.outputs) } - - /// Relevant (scripthash-wise) spent outpoints - fn spending(&self) -> impl Iterator + '_ { - self.spent.iter().copied() - } } // Confirmation height of a transaction or its mempool state: @@ -138,39 +133,38 @@ pub struct ScriptHashStatus { statushash: Option, // computed from history } -enum BalanceEntry { - Funded(OutPoint), - Spent(OutPoint), -} - /// Specific scripthash balance -#[derive(Default, Eq, PartialEq)] +#[derive(Default, Eq, PartialEq, Serialize)] pub(crate) struct Balance { - pub(crate) confirmed: Amount, - pub(crate) mempool_delta: SignedAmount, + #[serde(with = "bitcoin::util::amount::serde::as_sat")] + confirmed: Amount, + #[serde(with = "bitcoin::util::amount::serde::as_sat", rename = "unconfirmed")] + mempool_delta: SignedAmount, } #[derive(Default)] -struct Total { - funded: Amount, - spent: Amount, +struct Unspent { + map: HashMap, } -impl Total { +impl Unspent { fn balance(&self) -> Amount { - self.funded - self.spent + self.map.values().fold(Amount::default(), |acc, v| acc + *v) } - fn update( - &mut self, - entries: impl Iterator, - get_amount: impl Fn(OutPoint) -> Amount, - ) { - for entry in entries { - match entry { - BalanceEntry::Funded(outpoint) => self.funded += get_amount(outpoint), - BalanceEntry::Spent(outpoint) => self.spent += get_amount(outpoint), - } + fn insert(&mut self, entry: &TxEntry) { + for output in &entry.outputs { + let outpoint = OutPoint { + txid: entry.txid, + vout: output.index, + }; + self.map.insert(outpoint, output.value); + } + } + + fn remove(&mut self, entry: &TxEntry) { + for spent in &entry.spent { + self.map.remove(spent); } } } @@ -190,7 +184,7 @@ impl ScriptHashStatus { /// Iterate through confirmed TxEntries with their corresponding block heights. /// Skip entries from stale blocks. - fn confirmed_entries<'a>( + fn confirmed_height_entries<'a>( &'a self, chain: &'a Chain, ) -> impl Iterator + 'a { @@ -203,43 +197,36 @@ impl ScriptHashStatus { }) } + /// Iterate through confirmed TxEntries. + /// Skip entries from stale blocks. + fn confirmed_entries<'a>(&'a self, chain: &'a Chain) -> impl Iterator + 'a { + self.confirmed_height_entries(chain) + .flat_map(|(_height, entries)| entries) + } + /// Collect all funded and confirmed outpoints (as a set). fn confirmed_outpoints(&self, chain: &Chain) -> HashSet { self.confirmed_entries(chain) - .flat_map(|(_height, entries)| entries.iter().flat_map(TxEntry::funding)) + .flat_map(TxEntry::funding_outpoints) .collect() } - pub(crate) fn get_balance(&self, chain: &Chain, get_amount: F) -> Balance - where - F: Fn(OutPoint) -> Amount, - { - fn to_balance_entries<'a>( - entries: impl Iterator + 'a, - ) -> impl Iterator + 'a { - entries.flat_map(|e| { - let funded = e.funding().map(BalanceEntry::Funded); - let spent = e.spending().map(BalanceEntry::Spent); - funded.chain(spent) - }) - } + pub(crate) fn get_balance(&self, chain: &Chain) -> Balance { + let mut unspent = Unspent::default(); - let confirmed_entries = to_balance_entries( - self.confirmed_entries(chain) - .flat_map(|(_height, entries)| entries), - ); - let mempool_entries = to_balance_entries(self.mempool.iter()); + self.confirmed_entries(chain) + .for_each(|e| unspent.insert(e)); + self.confirmed_entries(chain) + .for_each(|e| unspent.remove(e)); + let confirmed_balance = unspent.balance(); - let mut total = Total::default(); - total.update(confirmed_entries, &get_amount); - let confirmed = total.balance(); - - total.update(mempool_entries, &get_amount); - let with_mempool = total.balance(); + self.mempool.iter().for_each(|e| unspent.insert(e)); + self.mempool.iter().for_each(|e| unspent.remove(e)); Balance { - confirmed, - mempool_delta: with_mempool.to_signed().unwrap() - confirmed.to_signed().unwrap(), + confirmed: confirmed_balance, + mempool_delta: unspent.balance().to_signed().unwrap() + - confirmed_balance.to_signed().unwrap(), } } @@ -251,7 +238,7 @@ impl ScriptHashStatus { /// Collect all confirmed history entries (in block order). fn get_confirmed_history(&self, chain: &Chain) -> Vec { - self.confirmed_entries(chain) + self.confirmed_height_entries(chain) .collect::>() .into_iter() .flat_map(|(height, entries)| { diff --git a/src/tracker.rs b/src/tracker.rs index 7bdaedd..f96fc0a 100644 --- a/src/tracker.rs +++ b/src/tracker.rs @@ -1,7 +1,6 @@ use anyhow::{Context, Result}; -use bitcoin::{BlockHash, OutPoint, Txid}; +use bitcoin::{BlockHash, Txid}; -use std::convert::TryFrom; use std::path::Path; use crate::{ @@ -76,17 +75,8 @@ impl Tracker { Ok(prev_statushash != status.statushash()) } - pub(crate) fn get_balance(&self, status: &ScriptHashStatus, cache: &Cache) -> Balance { - let get_amount_fn = |outpoint: OutPoint| { - cache - .get_tx(&outpoint.txid, |tx| { - let vout = usize::try_from(outpoint.vout).unwrap(); - bitcoin::Amount::from_sat(tx.output[vout].value) - }) - .expect("missing tx") - }; - - status.get_balance(self.chain(), get_amount_fn) + pub(crate) fn get_balance(&self, status: &ScriptHashStatus) -> Balance { + status.get_balance(self.chain()) } pub fn get_blockhash_by_txid(&self, txid: Txid) -> Option {