Add canonicalization modifier: assume not canonical.

crates/bitcoind_rpc/tests/test_emitter.rs

@@ -5,7 +5,7 @@ use bdk_chain::{
     bitcoin::{Address, Amount, Txid},
     local_chain::{CheckPoint, LocalChain},
     spk_txout::SpkTxOutIndex,
-    Balance, BlockId, IndexedTxGraph, Merge,
+    Balance, BlockId, CanonicalizationMods, IndexedTxGraph, Merge,
 };
 use bdk_testenv::{anyhow, TestEnv};
 use bitcoin::{hashes::Hash, Block, OutPoint, ScriptBuf, WScriptHash};
@@ -306,9 +306,13 @@ fn get_balance(
 ) -> anyhow::Result<Balance> {
     let chain_tip = recv_chain.tip().block_id();
     let outpoints = recv_graph.index.outpoints().clone();
-    let balance = recv_graph
-        .graph()
-        .balance(recv_chain, chain_tip, outpoints, |_, _| true);
+    let balance = recv_graph.graph().balance(
+        recv_chain,
+        chain_tip,
+        CanonicalizationMods::NONE,
+        outpoints,
+        |_, _| true,
+    );
     Ok(balance)
 }
 

crates/chain/benches/canonicalization.rs

@@ -1,3 +1,4 @@
+use bdk_chain::CanonicalizationMods;
 use bdk_chain::{keychain_txout::KeychainTxOutIndex, local_chain::LocalChain, IndexedTxGraph};
 use bdk_core::{BlockId, CheckPoint};
 use bdk_core::{ConfirmationBlockTime, TxUpdate};
@@ -90,16 +91,19 @@ fn setup<F: Fn(&mut KeychainTxGraph, &LocalChain)>(f: F) -> (KeychainTxGraph, Lo
 }
 
 fn run_list_canonical_txs(tx_graph: &KeychainTxGraph, chain: &LocalChain, exp_txs: usize) {
-    let txs = tx_graph
-        .graph()
-        .list_canonical_txs(chain, chain.tip().block_id());
+    let txs = tx_graph.graph().list_canonical_txs(
+        chain,
+        chain.tip().block_id(),
+        CanonicalizationMods::NONE,
+    );
     assert_eq!(txs.count(), exp_txs);
 }
 
 fn run_filter_chain_txouts(tx_graph: &KeychainTxGraph, chain: &LocalChain, exp_txos: usize) {
     let utxos = tx_graph.graph().filter_chain_txouts(
         chain,
         chain.tip().block_id(),
+        CanonicalizationMods::NONE,
         tx_graph.index.outpoints().clone(),
     );
     assert_eq!(utxos.count(), exp_txos);
@@ -109,6 +113,7 @@ fn run_filter_chain_unspents(tx_graph: &KeychainTxGraph, chain: &LocalChain, exp
     let utxos = tx_graph.graph().filter_chain_unspents(
         chain,
         chain.tip().block_id(),
+        CanonicalizationMods::NONE,
         tx_graph.index.outpoints().clone(),
     );
     assert_eq!(utxos.count(), exp_utxos);

crates/chain/src/canonical_iter.rs

@@ -4,19 +4,51 @@ use crate::{Anchor, ChainOracle, 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};
 
+/// Modifies the canonicalization algorithm.
+#[non_exhaustive]
+#[derive(Debug, Default, Clone)]
+pub struct CanonicalizationMods {
+    /// Transactions that will supercede all other transactions.
+    ///
+    /// In case of conflicting transactions within `assume_canonical`, transactions that appear
+    /// later in the list (have higher index) have precedence.
+    ///
+    /// If the same transaction exists in both `assume_canonical` and `assume_not_canonical`,
+    /// `assume_not_canonical` will take precedence.
+    pub assume_canonical: Vec<Txid>,
+
+    /// Transactions that will never be considered canonical.
+    ///
+    /// Descendants of these transactions will also be evicted.
+    ///
+    /// If the same transaction exists in both `assume_canonical` and `assume_not_canonical`,
+    /// `assume_not_canonical` will take precedence.
+    pub assume_not_canonical: Vec<Txid>,
+}
+
+impl CanonicalizationMods {
+    /// No mods.
+    pub const NONE: Self = Self {
+        assume_canonical: Vec::new(),
+        assume_not_canonical: Vec::new(),
+    };
+}
+
 /// Iterates over canonical txs.
 pub struct CanonicalIter<'g, A, C> {
     tx_graph: &'g TxGraph<A>,
     chain: &'g C,
     chain_tip: BlockId,
 
-    unprocessed_txs_with_anchors:
+    unprocessed_assumed_txs: Box<dyn Iterator<Item = (Txid, Arc<Transaction>)> + 'g>,
+    unprocessed_anchored_txs:
         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)>,
+    unprocessed_seen_txs: Box<dyn Iterator<Item = (Txid, Arc<Transaction>, u64)> + 'g>,
+    unprocessed_leftover_txs: VecDeque<(Txid, Arc<Transaction>, u32)>,
 
     canonical: HashMap<Txid, (Arc<Transaction>, CanonicalReason<A>)>,
     not_canonical: HashSet<Txid>,
@@ -26,14 +58,29 @@ pub struct CanonicalIter<'g, A, C> {
 
 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 {
+    pub fn new(
+        tx_graph: &'g TxGraph<A>,
+        chain: &'g C,
+        chain_tip: BlockId,
+        mods: CanonicalizationMods,
+    ) -> Self {
+        let mut not_canonical = HashSet::new();
+        for txid in mods.assume_not_canonical {
+            Self::_mark_not_canonical(tx_graph, &mut not_canonical, txid);
+        }
         let anchors = tx_graph.all_anchors();
-        let pending_anchored = Box::new(
+        let unprocessed_assumed_txs = Box::new(
+            mods.assume_canonical
+                .into_iter()
+                .rev()
+                .filter_map(|txid| Some((txid, tx_graph.get_tx(txid)?))),
+        );
+        let unprocessed_anchored_txs = 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(
+        let unprocessed_seen_txs = Box::new(
             tx_graph
                 .txids_by_descending_last_seen()
                 .filter_map(|(last_seen, txid)| Some((txid, tx_graph.get_tx(txid)?, last_seen))),
@@ -42,11 +89,12 @@ impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
             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(),
+            unprocessed_assumed_txs,
+            unprocessed_anchored_txs,
+            unprocessed_seen_txs,
+            unprocessed_leftover_txs: VecDeque::new(),
             canonical: HashMap::new(),
-            not_canonical: HashSet::new(),
+            not_canonical,
             queue: VecDeque::new(),
         }
     }
@@ -73,7 +121,7 @@ impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
             }
         }
         // cannot determine
-        self.unprocessed_txs_left_over.push_back((
+        self.unprocessed_leftover_txs.push_back((
             txid,
             tx,
             anchors
@@ -111,18 +159,11 @@ impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
                 // Any conflicts with a canonical tx can be added to `not_canonical`. Descendants
                 // of `not_canonical` txs can also be added to `not_canonical`.
                 for (_, conflict_txid) in self.tx_graph.direct_conflicts(&tx) {
-                    TxDescendants::new_include_root(
+                    Self::_mark_not_canonical(
                         self.tx_graph,
+                        &mut self.not_canonical,
                         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);
@@ -131,6 +172,17 @@ impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
         )
         .run_until_finished()
     }
+
+    fn _mark_not_canonical(tx_graph: &TxGraph<A>, not_canonical: &mut HashSet<Txid>, txid: Txid) {
+        TxDescendants::new_include_root(tx_graph, txid, |_: usize, txid: Txid| -> Option<()> {
+            if not_canonical.insert(txid) {
+                Some(())
+            } else {
+                None
+            }
+        })
+        .run_until_finished();
+    }
 }
 
 impl<A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'_, A, C> {
@@ -147,7 +199,13 @@ impl<A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'_, A, C> {
                 return Some(Ok((txid, tx, reason)));
             }
 
-            if let Some((txid, tx, anchors)) = self.unprocessed_txs_with_anchors.next() {
+            if let Some((txid, tx)) = self.unprocessed_assumed_txs.next() {
+                if !self.is_canonicalized(txid) {
+                    self.mark_canonical(txid, tx, CanonicalReason::assumed());
+                }
+            }
+
+            if let Some((txid, tx, anchors)) = self.unprocessed_anchored_txs.next() {
                 if !self.is_canonicalized(txid) {
                     if let Err(err) = self.scan_anchors(txid, tx, anchors) {
                         return Some(Err(err));
@@ -156,15 +214,15 @@ impl<A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'_, A, C> {
                 continue;
             }
 
-            if let Some((txid, tx, last_seen)) = self.unprocessed_txs_with_last_seens.next() {
+            if let Some((txid, tx, last_seen)) = self.unprocessed_seen_txs.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 let Some((txid, tx, height)) = self.unprocessed_leftover_txs.pop_front() {
                 if !self.is_canonicalized(txid) {
                     let observed_in = ObservedIn::Block(height);
                     self.mark_canonical(txid, tx, CanonicalReason::from_observed_in(observed_in));
@@ -189,6 +247,12 @@ pub enum ObservedIn {
 /// The reason why a transaction is canonical.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum CanonicalReason<A> {
+    /// This transaction is explicitly assumed to be canonical by the caller, superceding all other
+    /// canonicalization rules.
+    Assumed {
+        /// Whether it is a descendant that is assumed to be canonical.
+        descendant: Option<Txid>,
+    },
     /// This transaction is anchored in the best chain by `A`, and therefore canonical.
     Anchor {
         /// The anchor that anchored the transaction in the chain.
@@ -207,6 +271,12 @@ pub enum CanonicalReason<A> {
 }
 
 impl<A: Clone> CanonicalReason<A> {
+    /// Constructs a [`CanonicalReason`] for a transaction that is assumed to supercede all other
+    /// transactions.
+    pub fn assumed() -> Self {
+        Self::Assumed { descendant: None }
+    }
+
     /// Constructs a [`CanonicalReason`] from an `anchor`.
     pub fn from_anchor(anchor: A) -> Self {
         Self::Anchor {
@@ -229,6 +299,9 @@ impl<A: Clone> CanonicalReason<A> {
     /// descendant, but is transitively relevant.
     pub fn to_transitive(&self, descendant: Txid) -> Self {
         match self {
+            CanonicalReason::Assumed { .. } => Self::Assumed {
+                descendant: Some(descendant),
+            },
             CanonicalReason::Anchor { anchor, .. } => Self::Anchor {
                 anchor: anchor.clone(),
                 descendant: Some(descendant),
@@ -244,6 +317,7 @@ impl<A: Clone> CanonicalReason<A> {
     /// descendant.
     pub fn descendant(&self) -> &Option<Txid> {
         match self {
+            CanonicalReason::Assumed { descendant, .. } => descendant,
             CanonicalReason::Anchor { descendant, .. } => descendant,
             CanonicalReason::ObservedIn { descendant, .. } => descendant,
         }