use brk_types::{AddrBytes, AddrMempoolStats, Transaction, Txid};
use derive_more::Deref;
use rustc_hash::{FxHashMap, FxHashSet};
pub type AddrStats = (AddrMempoolStats, FxHashSet<Txid>);
#[derive(Default, Deref)]
pub struct AddrTracker(FxHashMap<AddrBytes, AddrStats>);
impl AddrTracker {
pub fn add_tx(&mut self, tx: &Transaction, txid: &Txid) {
self.update(tx, txid, true);
}
pub fn remove_tx(&mut self, tx: &Transaction, txid: &Txid) {
self.update(tx, txid, false);
}
fn update(&mut self, tx: &Transaction, txid: &Txid, is_addition: bool) {
for txin in &tx.input {
let Some(prevout) = txin.prevout.as_ref() else {
continue;
};
let Some(bytes) = prevout.addr_bytes() else {
continue;
};
let (stats, txids) = self.0.entry(bytes).or_default();
if is_addition {
txids.insert(txid.clone());
stats.sending(prevout);
} else {
txids.remove(txid);
stats.sent(prevout);
}
stats.update_tx_count(txids.len() as u32);
}
for txout in &tx.output {
let Some(bytes) = txout.addr_bytes() else {
continue;
};
let (stats, txids) = self.0.entry(bytes).or_default();
if is_addition {
txids.insert(txid.clone());
stats.receiving(txout);
} else {
txids.remove(txid);
stats.received(txout);
}
stats.update_tx_count(txids.len() as u32);
}
}
}