feat(chain)!: Better canonicalization algorithm

This is an O(n) algorithm to determine the canonical set of txids.

Author: evanlinjin

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,307 @@
+use core::cmp::Reverse;
+
+use crate::collections::{btree_set, hash_map, BTreeSet, HashMap, HashSet, VecDeque};
+use crate::tx_graph::{CanonicalTx, TxAncestors, TxDescendants, TxNode};
+use crate::{Anchor, ChainOracle, ChainPosition, TxGraph};
+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>)>;
+
+type ToProcess = btree_set::IntoIter<Reverse<(LastSeen, Txid)>>;
+
+/// Iterates over canonical txs.
+pub struct CanonicalIter<'g, A, C> {
+    tx_graph: &'g TxGraph<A>,
+    chain: &'g C,
+    chain_tip: BlockId,
+
+    to_process: ToProcess,
+    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 to_process = tx_graph
+            .full_txs()
+            .filter_map(|tx_node| {
+                Some(Reverse((
+                    tx_graph.last_seen_in(tx_node.txid)?,
+                    tx_node.txid,
+                )))
+            })
+            .collect::<BTreeSet<_>>()
+            .into_iter();
+        Self {
+            tx_graph,
+            chain,
+            chain_tip,
+            to_process,
+            canonical: HashMap::default(),
+            not_canonical: HashSet::default(),
+            queue: VecDeque::default(),
+        }
+    }
+
+    fn canonicalize_by_traversing_backwards(
+        &mut self,
+        txid: Txid,
+        last_seen: LastSeen,
+    ) -> Result<(), C::Error> {
+        type TxWithId = (Txid, Arc<Transaction>);
+        let tx = match self.tx_graph.get_tx(txid) {
+            Some(tx) => tx,
+            None => return Ok(()),
+        };
+        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) {
+                    for anchor in anchors {
+                        let anchor_block = anchor.anchor_block();
+                        match self.chain.is_block_in_chain(anchor_block, self.chain_tip) {
+                            Ok(None) | Ok(Some(false)) => continue,
+                            Ok(Some(true)) => {
+                                let reason = CanonicalReason::from_anchor(anchor.clone());
+                                self.mark_canonical(tx, reason);
+                                return None;
+                            }
+                            Err(err) => return Some(Err(err)),
+                        };
+                    }
+                }
+                // Coinbase transactions cannot exist in mempool.
+                if tx.is_coinbase() {
+                    return None;
+                }
+                for (_, conflicting_txid) in self.tx_graph.direct_conflicts(&tx) {
+                    if self.canonical.contains_key(&conflicting_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
+            }
+        })
+        .for_each(|_| {})
+    }
+
+    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(())
+                    }
+                }
+            },
+        )
+        .for_each(|_| {})
+    }
+}
+
+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)));
+            }
+
+            let Reverse((last_seen, txid)) = self.to_process.next()?;
+            if let Err(err) = self.canonicalize_by_traversing_backwards(txid, last_seen) {
+                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 LastSeen {
+    /// The transaction was last seen in the mempool at the given unix timestamp.
+    Mempool(u64),
+    /// The transaction was last seen in a block of height.
+    Block(u32),
+}
+
+/// 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 [`LastSeen`] value of the transaction.
+        last_seen: LastSeen,
+        /// Whether the [`LastSeen`] 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: LastSeen) -> 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: LastSeen) -> Self {
+        Self::LastSeen {
+            last_seen,
+            descendant: Some(descendant),
+        }
+    }
+
+    /// Adds a `descendant` to the [`CanonicalReason`].
+    ///
+    /// This signals that either the [`LastSeen`] 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 [`LastSeen`] 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(desc_txid) => {
+                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),
+                    None => ChainPosition::ConfirmedByTransitivity(desc_txid, anchor),
+                }
+            }
+            None => ChainPosition::Confirmed(anchor),
+        },
+        CanonicalReason::LastSeen { last_seen, .. } => match last_seen {
+            LastSeen::Mempool(last_seen) => ChainPosition::Unconfirmed(last_seen),
+            LastSeen::Block(_) => ChainPosition::UnconfirmedAndNotSeen,
+        },
+    };
+    Ok(CanonicalTx {
+        chain_position,
+        tx_node,
+    })
+}

crates/chain/src/chain_data.rs

@@ -15,16 +15,24 @@ use crate::{Anchor, COINBASE_MATURITY};
     ))
 )]
 pub enum ChainPosition<A> {
-    /// The chain data is seen as confirmed, and in anchored by `A`.
+    /// The chain data is confirmed because it is anchored by `A`.
     Confirmed(A),
-    /// The chain data is not confirmed and last seen in the mempool at this timestamp.
+    /// The chain data is confirmed because it has a descendant that is anchored by `A`.
+    ConfirmedByTransitivity(Txid, A),
+    /// The chain data is not confirmed and is last seen in the mempool (or has a descendant that
+    /// is last seen in the mempool) at this timestamp.
     Unconfirmed(u64),
+    /// The chain data is not confirmed and we have never seen it in the mempool.
+    UnconfirmedAndNotSeen,
 }
 
 impl<A> ChainPosition<A> {
     /// Returns whether [`ChainPosition`] is confirmed or not.
     pub fn is_confirmed(&self) -> bool {
-        matches!(self, Self::Confirmed(_))
+        matches!(
+            self,
+            Self::Confirmed(_) | Self::ConfirmedByTransitivity(_, _)
+        )
     }
 }
 
@@ -33,7 +41,11 @@ impl<A: Clone> ChainPosition<&A> {
     pub fn cloned(self) -> ChainPosition<A> {
         match self {
             ChainPosition::Confirmed(a) => ChainPosition::Confirmed(a.clone()),
+            ChainPosition::ConfirmedByTransitivity(txid, a) => {
+                ChainPosition::ConfirmedByTransitivity(txid, a.clone())
+            }
             ChainPosition::Unconfirmed(last_seen) => ChainPosition::Unconfirmed(last_seen),
+            ChainPosition::UnconfirmedAndNotSeen => ChainPosition::UnconfirmedAndNotSeen,
         }
     }
 }
@@ -42,8 +54,10 @@ impl<A: Anchor> ChainPosition<A> {
     /// Determines the upper bound of the confirmation height.
     pub fn confirmation_height_upper_bound(&self) -> Option<u32> {
         match self {
-            ChainPosition::Confirmed(a) => Some(a.confirmation_height_upper_bound()),
-            ChainPosition::Unconfirmed(_) => None,
+            ChainPosition::Confirmed(a) | ChainPosition::ConfirmedByTransitivity(_, a) => {
+                Some(a.confirmation_height_upper_bound())
+            }
+            ChainPosition::Unconfirmed(_) | ChainPosition::UnconfirmedAndNotSeen => None,
         }
     }
 }
@@ -73,9 +87,10 @@ impl<A: Anchor> FullTxOut<A> {
     /// [`confirmation_height_upper_bound`]: Anchor::confirmation_height_upper_bound
     pub fn is_mature(&self, tip: u32) -> bool {
         if self.is_on_coinbase {
-            let tx_height = match &self.chain_position {
-                ChainPosition::Confirmed(anchor) => anchor.confirmation_height_upper_bound(),
-                ChainPosition::Unconfirmed(_) => {
+            let tx_height = self.chain_position.confirmation_height_upper_bound();
+            let tx_height = match tx_height {
+                Some(tx_height) => tx_height,
+                None => {
                     debug_assert!(false, "coinbase tx can never be unconfirmed");
                     return false;
                 }
@@ -103,9 +118,9 @@ impl<A: Anchor> FullTxOut<A> {
             return false;
         }
 
-        let confirmation_height = match &self.chain_position {
-            ChainPosition::Confirmed(anchor) => anchor.confirmation_height_upper_bound(),
-            ChainPosition::Unconfirmed(_) => return false,
+        let confirmation_height = match self.chain_position.confirmation_height_upper_bound() {
+            Some(h) => h,
+            None => return false,
         };
         if confirmation_height > tip {
             return false;