diff --git a/crates/chain/src/tx_graph.rs b/crates/chain/src/tx_graph.rs index 73d602df..4c0cdb3d 100644 --- a/crates/chain/src/tx_graph.rs +++ b/crates/chain/src/tx_graph.rs @@ -55,6 +55,7 @@ use crate::{ ChainOracle, ChainPosition, FullTxOut, }; use alloc::vec::Vec; +use alloc::collections::vec_deque::VecDeque; use bitcoin::{OutPoint, Script, Transaction, TxOut, Txid}; use core::{ convert::Infallible, @@ -319,6 +320,30 @@ impl TxGraph { .map(|(outpoint, spends)| (outpoint.vout, spends)) } + /// Creates an iterator that filters and maps ancestor transactions. + /// + /// The iterator starts with the ancestors of the supplied `tx` (ancestor transactions of `tx` + /// are transactions spent by `tx`). The supplied transaction is excluded from the iterator. + /// + /// The supplied closure takes in two inputs `(depth, ancestor_tx)`: + /// + /// * `depth` is the distance between the starting `Transaction` and the `ancestor_tx`. I.e., if + /// the `Transaction` is spending an output of the `ancestor_tx` then `depth` will be 1. + /// * `ancestor_tx` is the `Transaction`'s ancestor which we are considering to walk. + /// + /// The supplied closure returns an `Option`, allowing the caller to map each `Transaction` + /// it visits and decide whether to visit ancestors. + pub fn walk_ancestors<'g, F, O>( + &'g self, + tx: &'g Transaction, + walk_map: F, + ) -> TxAncestors<'g, A, F> + where + F: FnMut(usize, &'g Transaction) -> Option + 'g, + { + TxAncestors::new_exclude_root(self, tx, walk_map) + } + /// Creates an iterator that filters and maps descendants from the starting `txid`. /// /// The supplied closure takes in two inputs `(depth, descendant_txid)`: @@ -1137,6 +1162,126 @@ impl AsRef> for TxGraph { } } +/// An iterator that traverses ancestors of a given root transaction. +/// +/// The iterator excludes partial transactions. +/// +/// This `struct` is created by the [`walk_ancestors`] method of [`TxGraph`]. +/// +/// [`walk_ancestors`]: TxGraph::walk_ancestors +pub struct TxAncestors<'g, A, F> { + graph: &'g TxGraph, + visited: HashSet, + queue: VecDeque<(usize, &'g Transaction)>, + filter_map: F, +} + +impl<'g, A, F> TxAncestors<'g, A, F> { + /// Creates a `TxAncestors` that includes the starting `Transaction` when iterating. + pub(crate) fn new_include_root( + graph: &'g TxGraph, + tx: &'g Transaction, + filter_map: F, + ) -> Self { + Self { + graph, + visited: Default::default(), + queue: [(0, tx)].into(), + filter_map, + } + } + + /// Creates a `TxAncestors` that excludes the starting `Transaction` when iterating. + pub(crate) fn new_exclude_root( + graph: &'g TxGraph, + tx: &'g Transaction, + filter_map: F, + ) -> Self { + let mut ancestors = Self { + graph, + visited: Default::default(), + queue: Default::default(), + filter_map, + }; + ancestors.populate_queue(1, tx); + ancestors + } + + /// Creates a `TxAncestors` from multiple starting `Transaction`s that includes the starting + /// `Transaction`s when iterating. + #[allow(unused)] + pub(crate) fn from_multiple_include_root( + graph: &'g TxGraph, + txs: I, + filter_map: F, + ) -> Self + where + I: IntoIterator, + { + Self { + graph, + visited: Default::default(), + queue: txs.into_iter().map(|tx| (0, tx)).collect(), + filter_map, + } + } + + /// Creates a `TxAncestors` from multiple starting `Transaction`s that excludes the starting + /// `Transaction`s when iterating. + #[allow(unused)] + pub(crate) fn from_multiple_exclude_root( + graph: &'g TxGraph, + txs: I, + filter_map: F, + ) -> Self + where + I: IntoIterator, + { + let mut ancestors = Self { + graph, + visited: Default::default(), + queue: Default::default(), + filter_map, + }; + for tx in txs { + ancestors.populate_queue(1, tx); + } + ancestors + } + + fn populate_queue(&mut self, depth: usize, tx: &'g Transaction) { + let ancestors = tx + .input + .iter() + .map(|txin| txin.previous_output.txid) + .filter(|&prev_txid| self.visited.insert(prev_txid)) + .filter_map(|prev_txid| self.graph.get_tx(prev_txid)) + .map(|tx| (depth, tx)); + self.queue.extend(ancestors); + } +} + +impl<'g, A, F, O> Iterator for TxAncestors<'g, A, F> +where + F: FnMut(usize, &'g Transaction) -> Option, +{ + type Item = O; + + fn next(&mut self) -> Option { + loop { + // we have exhausted all paths when queue is empty + let (ancestor_depth, tx) = self.queue.pop_front()?; + // ignore paths when user filters them out + let item = match (self.filter_map)(ancestor_depth, tx) { + Some(item) => item, + None => continue, + }; + self.populate_queue(ancestor_depth + 1, tx); + return Some(item); + } + } +} + /// An iterator that traverses transaction descendants. /// /// This `struct` is created by the [`walk_descendants`] method of [`TxGraph`].