use crate::collections::{hash_map, HashMap, HashSet, VecDeque};
use crate::tx_graph::{TxAncestors, TxDescendants};
use crate::{Anchor, ChainOracle, TxGraph};
use alloc::boxed::Box;
use alloc::collections::BTreeSet;
use alloc::sync::Arc;
use bdk_core::BlockId;
use bitcoin::{Transaction, Txid};
pub struct CanonicalIter<'g, A, C> {
tx_graph: &'g TxGraph<A>,
chain: &'g C,
chain_tip: BlockId,
unprocessed_txs_with_anchors:
Box<dyn Iterator<Item = (Txid, Arc<Transaction>, &'g BTreeSet<A>)> + 'g>,
unprocessed_txs_with_last_seens: Box<dyn Iterator<Item = (Txid, Arc<Transaction>, u64)> + 'g>,
unprocessed_txs_left_over: VecDeque<(Txid, Arc<Transaction>, u32)>,
canonical: HashMap<Txid, (Arc<Transaction>, CanonicalReason<A>)>,
not_canonical: HashSet<Txid>,
queue: VecDeque<Txid>,
}
impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
pub fn new(tx_graph: &'g TxGraph<A>, chain: &'g C, chain_tip: BlockId) -> Self {
let anchors = tx_graph.all_anchors();
let pending_anchored = Box::new(
tx_graph
.txids_by_descending_anchor_height()
.filter_map(|(_, txid)| Some((txid, tx_graph.get_tx(txid)?, anchors.get(&txid)?))),
);
let pending_last_seen = Box::new(
tx_graph
.txids_by_descending_last_seen()
.filter_map(|(last_seen, txid)| Some((txid, tx_graph.get_tx(txid)?, last_seen))),
);
Self {
tx_graph,
chain,
chain_tip,
unprocessed_txs_with_anchors: pending_anchored,
unprocessed_txs_with_last_seens: pending_last_seen,
unprocessed_txs_left_over: VecDeque::new(),
canonical: HashMap::new(),
not_canonical: HashSet::new(),
queue: VecDeque::new(),
}
}
fn is_canonicalized(&self, txid: Txid) -> bool {
self.canonical.contains_key(&txid) || self.not_canonical.contains(&txid)
}
fn scan_anchors(
&mut self,
txid: Txid,
tx: Arc<Transaction>,
anchors: &BTreeSet<A>,
) -> Result<(), C::Error> {
for anchor in anchors {
let in_chain_opt = self
.chain
.is_block_in_chain(anchor.anchor_block(), self.chain_tip)?;
if in_chain_opt == Some(true) {
self.mark_canonical(txid, tx, CanonicalReason::from_anchor(anchor.clone()));
return Ok(());
}
}
self.unprocessed_txs_left_over.push_back((
txid,
tx,
anchors
.iter()
.last()
.expect(
"tx taken from `unprocessed_txs_with_anchors` so it must atleast have an anchor",
)
.confirmation_height_upper_bound(),
));
Ok(())
}
fn mark_canonical(&mut self, txid: Txid, tx: Arc<Transaction>, reason: CanonicalReason<A>) {
let starting_txid = txid;
let mut is_root = true;
TxAncestors::new_include_root(
self.tx_graph,
tx,
|_: usize, tx: Arc<Transaction>| -> Option<()> {
let this_txid = tx.compute_txid();
let this_reason = if is_root {
is_root = false;
reason.clone()
} else {
reason.to_transitive(starting_txid)
};
let canonical_entry = match self.canonical.entry(this_txid) {
hash_map::Entry::Occupied(_) => return None,
hash_map::Entry::Vacant(entry) => entry,
};
for (_, conflict_txid) in self.tx_graph.direct_conflicts(&tx) {
TxDescendants::new_include_root(
self.tx_graph,
conflict_txid,
|_: usize, txid: Txid| -> Option<()> {
if self.not_canonical.insert(txid) {
Some(())
} else {
None
}
},
)
.run_until_finished()
}
canonical_entry.insert((tx, this_reason));
self.queue.push_back(this_txid);
Some(())
},
)
.run_until_finished()
}
}
impl<A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'_, A, C> {
type Item = Result<(Txid, Arc<Transaction>, CanonicalReason<A>), C::Error>;
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(txid) = self.queue.pop_front() {
let (tx, reason) = self
.canonical
.get(&txid)
.cloned()
.expect("reason must exist");
return Some(Ok((txid, tx, reason)));
}
if let Some((txid, tx, anchors)) = self.unprocessed_txs_with_anchors.next() {
if !self.is_canonicalized(txid) {
if let Err(err) = self.scan_anchors(txid, tx, anchors) {
return Some(Err(err));
}
}
continue;
}
if let Some((txid, tx, last_seen)) = self.unprocessed_txs_with_last_seens.next() {
if !self.is_canonicalized(txid) {
let observed_in = ObservedIn::Mempool(last_seen);
self.mark_canonical(txid, tx, CanonicalReason::from_observed_in(observed_in));
}
continue;
}
if let Some((txid, tx, height)) = self.unprocessed_txs_left_over.pop_front() {
if !self.is_canonicalized(txid) {
let observed_in = ObservedIn::Block(height);
self.mark_canonical(txid, tx, CanonicalReason::from_observed_in(observed_in));
}
continue;
}
return None;
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ObservedIn {
Block(u32),
Mempool(u64),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CanonicalReason<A> {
Anchor {
anchor: A,
descendant: Option<Txid>,
},
ObservedIn {
observed_in: ObservedIn,
descendant: Option<Txid>,
},
}
impl<A: Clone> CanonicalReason<A> {
pub fn from_anchor(anchor: A) -> Self {
Self::Anchor {
anchor,
descendant: None,
}
}
pub fn from_observed_in(observed_in: ObservedIn) -> Self {
Self::ObservedIn {
observed_in,
descendant: None,
}
}
pub fn to_transitive(&self, descendant: Txid) -> Self {
match self {
CanonicalReason::Anchor { anchor, .. } => Self::Anchor {
anchor: anchor.clone(),
descendant: Some(descendant),
},
CanonicalReason::ObservedIn { observed_in, .. } => Self::ObservedIn {
observed_in: *observed_in,
descendant: Some(descendant),
},
}
}
pub fn descendant(&self) -> &Option<Txid> {
match self {
CanonicalReason::Anchor { descendant, .. } => descendant,
CanonicalReason::ObservedIn { descendant, .. } => descendant,
}
}
}