Introduce O(n) canonicalization algorithm

crates/bitcoind_rpc/tests/test_emitter.rs

@@ -389,6 +389,7 @@ fn tx_can_become_unconfirmed_after_reorg() -> anyhow::Result<()> {
         assert_eq!(
             get_balance(&recv_chain, &recv_graph)?,
             Balance {
+                trusted_pending: SEND_AMOUNT * reorg_count as u64,
                 confirmed: SEND_AMOUNT * (ADDITIONAL_COUNT - reorg_count) as u64,
                 ..Balance::default()
             },

crates/chain/src/canonical_iter.rs

@@ -0,0 +1,341 @@
+use crate::collections::{hash_map, HashMap, HashSet, VecDeque};
+use crate::tx_graph::{CanonicalTx, TxAncestors, TxDescendants, TxNode};
+use crate::{Anchor, ChainOracle, ChainPosition, TxGraph};
+use alloc::boxed::Box;
+use alloc::collections::BTreeSet;
+use alloc::sync::Arc;
+use alloc::vec::Vec;
+use bdk_core::BlockId;
+use bitcoin::{Transaction, Txid};
+
+/// A set of canonical transactions.
+pub type CanonicalSet<A> = HashMap<Txid, (Arc<Transaction>, CanonicalReason<A>)>;
+
+/// Iterates over canonical txs.
+pub struct CanonicalIter<'g, A, C> {
+    tx_graph: &'g TxGraph<A>,
+    chain: &'g C,
+    chain_tip: BlockId,
+
+    todo_anchored: Box<dyn Iterator<Item = (Txid, Arc<Transaction>, &'g BTreeSet<A>)> + 'g>,
+    todo_last_seen: Box<dyn Iterator<Item = (Txid, Arc<Transaction>, LastSeenIn)> + 'g>,
+
+    canonical: CanonicalSet<A>,
+    not_canonical: HashSet<Txid>,
+
+    queue: VecDeque<Txid>,
+}
+
+impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
+    /// Constructs [`CanonicalIter`].
+    pub fn new(tx_graph: &'g TxGraph<A>, chain: &'g C, chain_tip: BlockId) -> Self {
+        let todo_anchored = Box::new(
+            tx_graph
+                .all_anchors()
+                .iter()
+                .filter_map(|(&txid, a)| Some((txid, tx_graph.get_tx(txid)?, a))),
+        );
+        let todo_last_seen = Box::new(
+            tx_graph
+                .txs_by_last_seen_in()
+                .filter_map(|(lsi, txid)| Some((txid, tx_graph.get_tx(txid)?, lsi))),
+        );
+        Self {
+            tx_graph,
+            chain,
+            chain_tip,
+            todo_anchored,
+            todo_last_seen,
+            canonical: HashMap::default(),
+            not_canonical: HashSet::default(),
+            queue: VecDeque::default(),
+        }
+    }
+
+    fn is_canonical_by_anchor(
+        &mut self,
+        tx: Arc<Transaction>,
+        anchors: &BTreeSet<A>,
+    ) -> Result<bool, 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) {
+                for (_, conflict_txid) in self.tx_graph.direct_conflicts(&tx) {
+                    self.mark_not_canonical(conflict_txid);
+                }
+                self.mark_canonical(tx, CanonicalReason::from_anchor(anchor.clone()));
+                return Ok(true);
+            }
+        }
+        Ok(false)
+    }
+
+    fn canonicalize_by_traversing_backwards(
+        &mut self,
+        txid: Txid,
+        tx: Arc<Transaction>,
+        last_seen: LastSeenIn,
+    ) -> Result<(), C::Error> {
+        type TxWithId = (Txid, Arc<Transaction>);
+        let maybe_canonical = TxAncestors::new_include_root(
+            self.tx_graph,
+            tx,
+            |_: usize, tx: Arc<Transaction>| -> Option<Result<TxWithId, C::Error>> {
+                let txid = tx.compute_txid();
+                if self.canonical.contains_key(&txid) || self.not_canonical.contains(&txid) {
+                    return None;
+                }
+                if let Some(anchors) = self.tx_graph.all_anchors().get(&txid) {
+                    match self.is_canonical_by_anchor(tx.clone(), anchors) {
+                        Ok(false) => {}
+                        Ok(true) => return None,
+                        Err(err) => return Some(Err(err)),
+                    };
+                }
+                // Coinbase transactions cannot exist in mempool.
+                if tx.is_coinbase() {
+                    return None;
+                }
+                for (_, conflict_txid) in self.tx_graph.direct_conflicts(&tx) {
+                    if self.canonical.contains_key(&conflict_txid) {
+                        self.mark_not_canonical(txid);
+                        return None;
+                    }
+                }
+                Some(Ok((txid, tx)))
+            },
+        )
+        .collect::<Result<Vec<_>, C::Error>>()?;
+
+        // This assumes that `last_seen` values are fully transitive. I.e. if A is an ancestor of B,
+        // then the most recent timestamp between A & B also applies to A.
+        let starting_txid = txid;
+        for (txid, tx) in maybe_canonical {
+            if self.not_canonical.contains(&txid) {
+                continue;
+            }
+            self.mark_canonical(
+                tx,
+                if txid == starting_txid {
+                    CanonicalReason::<A>::from_last_seen(last_seen)
+                } else {
+                    CanonicalReason::<A>::from_descendant_last_seen(starting_txid, last_seen)
+                },
+            );
+        }
+        Ok(())
+    }
+
+    fn mark_not_canonical(&mut self, txid: Txid) {
+        TxDescendants::new_include_root(self.tx_graph, txid, |_: usize, txid: Txid| -> Option<()> {
+            if self.not_canonical.insert(txid) {
+                Some(())
+            } else {
+                None
+            }
+        })
+        .run_until_finished()
+    }
+
+    fn mark_canonical(&mut self, tx: Arc<Transaction>, reason: CanonicalReason<A>) {
+        let starting_txid = tx.compute_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.clone().with_descendant(starting_txid)
+                };
+                match self.canonical.entry(this_txid) {
+                    hash_map::Entry::Occupied(_) => None,
+                    hash_map::Entry::Vacant(entry) => {
+                        entry.insert((tx, this_reason));
+                        self.queue.push_back(this_txid);
+                        Some(())
+                    }
+                }
+            },
+        )
+        .run_until_finished()
+    }
+}
+
+impl<'g, A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'g, 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.todo_anchored.next() {
+                if self.canonical.contains_key(&txid) || self.not_canonical.contains(&txid) {
+                    continue;
+                }
+                if let Err(err) = self.is_canonical_by_anchor(tx, anchors) {
+                    return Some(Err(err));
+                }
+                continue;
+            }
+
+            let (txid, tx, lsi) = self.todo_last_seen.next()?;
+            if let Err(err) = self.canonicalize_by_traversing_backwards(txid, tx, lsi) {
+                return Some(Err(err));
+            }
+        }
+    }
+}
+
+/// Represents when and where a given transaction is last seen.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, core::hash::Hash)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub enum LastSeenIn {
+    /// The transaction was last seen in a block of height.
+    Block(u32),
+    /// The transaction was last seen in the mempool at the given unix timestamp.
+    Mempool(u64),
+}
+
+/// The reason why a transaction is canonical.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum CanonicalReason<A> {
+    /// This transaction is anchored in the best chain by `A`, and therefore canonical.
+    Anchor {
+        /// The anchor that anchored the transaction in the chain.
+        anchor: A,
+        /// Whether the anchor is of the transaction's descendant.
+        descendant: Option<Txid>,
+    },
+    /// This transaction does not conflict with any other transaction with a more recent `last_seen`
+    /// value or one that is anchored in the best chain.
+    LastSeen {
+        /// The [`LastSeenIn`] value of the transaction.
+        last_seen: LastSeenIn,
+        /// Whether the [`LastSeenIn`] value is of the transaction's descendant.
+        descendant: Option<Txid>,
+    },
+}
+
+impl<A> CanonicalReason<A> {
+    /// Constructs a [`CanonicalReason`] from an `anchor`.
+    pub fn from_anchor(anchor: A) -> Self {
+        Self::Anchor {
+            anchor,
+            descendant: None,
+        }
+    }
+
+    /// Constructs a [`CanonicalReason`] from a `descendant`'s `anchor`.
+    pub fn from_descendant_anchor(descendant: Txid, anchor: A) -> Self {
+        Self::Anchor {
+            anchor,
+            descendant: Some(descendant),
+        }
+    }
+
+    /// Constructs a [`CanonicalReason`] from a `last_seen` value.
+    pub fn from_last_seen(last_seen: LastSeenIn) -> Self {
+        Self::LastSeen {
+            last_seen,
+            descendant: None,
+        }
+    }
+
+    /// Constructs a [`CanonicalReason`] from a `descendant`'s `last_seen` value.
+    pub fn from_descendant_last_seen(descendant: Txid, last_seen: LastSeenIn) -> Self {
+        Self::LastSeen {
+            last_seen,
+            descendant: Some(descendant),
+        }
+    }
+
+    /// Adds a `descendant` to the [`CanonicalReason`].
+    ///
+    /// This signals that either the [`LastSeenIn`] or [`Anchor`] value belongs to the transaction's
+    /// descendant.
+    #[must_use]
+    pub fn with_descendant(self, descendant: Txid) -> Self {
+        match self {
+            CanonicalReason::Anchor { anchor, .. } => Self::Anchor {
+                anchor,
+                descendant: Some(descendant),
+            },
+            CanonicalReason::LastSeen { last_seen, .. } => Self::LastSeen {
+                last_seen,
+                descendant: Some(descendant),
+            },
+        }
+    }
+
+    /// This signals that either the [`LastSeenIn`] or [`Anchor`] value belongs to the transaction's
+    /// descendant.
+    pub fn descendant(&self) -> &Option<Txid> {
+        match self {
+            CanonicalReason::Anchor { descendant, .. } => descendant,
+            CanonicalReason::LastSeen { descendant, .. } => descendant,
+        }
+    }
+}
+
+/// Helper to create canonical tx.
+pub fn make_canonical_tx<'a, A: Anchor, C: ChainOracle>(
+    chain: &C,
+    chain_tip: BlockId,
+    tx_node: TxNode<'a, Arc<Transaction>, A>,
+    canonical_reason: CanonicalReason<A>,
+) -> Result<CanonicalTx<'a, Arc<Transaction>, A>, C::Error> {
+    let chain_position = match canonical_reason {
+        CanonicalReason::Anchor { anchor, descendant } => match descendant {
+            Some(_) => {
+                let direct_anchor = tx_node
+                    .anchors
+                    .iter()
+                    .find_map(|a| -> Option<Result<A, C::Error>> {
+                        match chain.is_block_in_chain(a.anchor_block(), chain_tip) {
+                            Ok(Some(true)) => Some(Ok(a.clone())),
+                            Ok(Some(false)) | Ok(None) => None,
+                            Err(err) => Some(Err(err)),
+                        }
+                    })
+                    .transpose()?;
+                match direct_anchor {
+                    Some(anchor) => ChainPosition::Confirmed {
+                        anchor,
+                        transitively: None,
+                    },
+                    None => ChainPosition::Confirmed {
+                        anchor,
+                        transitively: descendant,
+                    },
+                }
+            }
+            None => ChainPosition::Confirmed {
+                anchor,
+                transitively: None,
+            },
+        },
+        CanonicalReason::LastSeen { last_seen, .. } => match last_seen {
+            LastSeenIn::Mempool(last_seen) => ChainPosition::Unconfirmed {
+                last_seen: Some(last_seen),
+            },
+            LastSeenIn::Block(_) => ChainPosition::Unconfirmed { last_seen: None },
+        },
+    };
+    Ok(CanonicalTx {
+        chain_position,
+        tx_node,
+    })
+}