feat(chain,wallet)!: Add ability to modify canonicalization algorithm

Introduce `CanonicalizationMods` which is passed in to
`CanonicalIter::new`.

`CanonicalizationMods::assume_canonical` is the only field right now.
This contains a list of txids that we assume to be canonical,
superceding any other canonicalization rules.

Author: evanlinjin

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};
@@ -92,14 +93,15 @@ 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());
+        .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 +111,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,15 +4,34 @@ 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.
+#[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.
+    pub assume_canonical: Vec<Txid>,
+}
+
+impl CanonicalizationMods {
+    /// No mods.
+    pub const NONE: Self = Self {
+        assume_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_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_seen_txs: Box<dyn Iterator<Item = (Txid, Arc<Transaction>, u64)> + 'g>,
@@ -26,8 +45,19 @@ 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 anchors = tx_graph.all_anchors();
+        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()
@@ -42,6 +72,7 @@ impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
             tx_graph,
             chain,
             chain_tip,
+            unprocessed_assumed_txs,
             unprocessed_anchored_txs,
             unprocessed_seen_txs,
             unprocessed_leftover_txs: VecDeque::new(),
@@ -147,6 +178,12 @@ impl<A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'_, A, C> {
                 return Some(Ok((txid, tx, reason)));
             }
 
+            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) {
@@ -189,6 +226,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 +250,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 +278,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 +296,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,
         }

crates/chain/src/tx_graph.rs

@@ -94,6 +94,7 @@ use crate::collections::*;
 use crate::BlockId;
 use crate::CanonicalIter;
 use crate::CanonicalReason;
+use crate::CanonicalizationMods;
 use crate::ObservedIn;
 use crate::{Anchor, Balance, ChainOracle, ChainPosition, FullTxOut, Merge};
 use alloc::collections::vec_deque::VecDeque;
@@ -829,25 +830,46 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
     ) -> impl Iterator<Item = Result<CanonicalTx<'a, Arc<Transaction>, A>, C::Error>> {
-        self.canonical_iter(chain, chain_tip).flat_map(move |res| {
-            res.map(|(txid, _, canonical_reason)| {
-                let tx_node = self.get_tx_node(txid).expect("must contain tx");
-                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 {
+        fn find_direct_anchor<A: Anchor, C: ChainOracle>(
+            tx_node: &TxNode<'_, Arc<Transaction>, A>,
+            chain: &C,
+            chain_tip: BlockId,
+        ) -> Result<Option<A>, C::Error> {
+            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()
+        }
+        self.canonical_iter(chain, chain_tip, mods)
+            .flat_map(move |res| {
+                res.map(|(txid, _, canonical_reason)| {
+                    let tx_node = self.get_tx_node(txid).expect("must contain tx");
+                    let chain_position = match canonical_reason {
+                        CanonicalReason::Assumed { descendant } => match descendant {
+                            Some(_) => match find_direct_anchor(&tx_node, chain, chain_tip)? {
+                                Some(anchor) => ChainPosition::Confirmed {
+                                    anchor,
+                                    transitively: None,
+                                },
+                                None => ChainPosition::Unconfirmed {
+                                    last_seen: tx_node.last_seen_unconfirmed,
+                                },
+                            },
+                            None => ChainPosition::Unconfirmed {
+                                last_seen: tx_node.last_seen_unconfirmed,
+                            },
+                        },
+                        CanonicalReason::Anchor { anchor, descendant } => match descendant {
+                            Some(_) => match find_direct_anchor(&tx_node, chain, chain_tip)? {
                                 Some(anchor) => ChainPosition::Confirmed {
                                     anchor,
                                     transitively: None,
@@ -856,26 +878,25 @@ impl<A: Anchor> TxGraph<A> {
                                     anchor,
                                     transitively: descendant,
                                 },
-                            }
-                        }
-                        None => ChainPosition::Confirmed {
-                            anchor,
-                            transitively: None,
+                            },
+                            None => ChainPosition::Confirmed {
+                                anchor,
+                                transitively: None,
+                            },
                         },
-                    },
-                    CanonicalReason::ObservedIn { observed_in, .. } => match observed_in {
-                        ObservedIn::Mempool(last_seen) => ChainPosition::Unconfirmed {
-                            last_seen: Some(last_seen),
+                        CanonicalReason::ObservedIn { observed_in, .. } => match observed_in {
+                            ObservedIn::Mempool(last_seen) => ChainPosition::Unconfirmed {
+                                last_seen: Some(last_seen),
+                            },
+                            ObservedIn::Block(_) => ChainPosition::Unconfirmed { last_seen: None },
                         },
-                        ObservedIn::Block(_) => ChainPosition::Unconfirmed { last_seen: None },
-                    },
-                };
-                Ok(CanonicalTx {
-                    chain_position,
-                    tx_node,
+                    };
+                    Ok(CanonicalTx {
+                        chain_position,
+                        tx_node,
+                    })
                 })
             })
-        })
     }
 
     /// List graph transactions that are in `chain` with `chain_tip`.
@@ -887,8 +908,9 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
     ) -> impl Iterator<Item = CanonicalTx<'a, Arc<Transaction>, A>> {
-        self.try_list_canonical_txs(chain, chain_tip)
+        self.try_list_canonical_txs(chain, chain_tip, mods)
             .map(|res| res.expect("infallible"))
     }
 
@@ -915,11 +937,12 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
         outpoints: impl IntoIterator<Item = (OI, OutPoint)> + 'a,
     ) -> Result<impl Iterator<Item = (OI, FullTxOut<A>)> + 'a, C::Error> {
         let mut canon_txs = HashMap::<Txid, CanonicalTx<Arc<Transaction>, A>>::new();
         let mut canon_spends = HashMap::<OutPoint, Txid>::new();
-        for r in self.try_list_canonical_txs(chain, chain_tip) {
+        for r in self.try_list_canonical_txs(chain, chain_tip, mods) {
             let canonical_tx = r?;
             let txid = canonical_tx.tx_node.txid;
 
@@ -988,8 +1011,9 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
     ) -> CanonicalIter<'a, A, C> {
-        CanonicalIter::new(self, chain, chain_tip)
+        CanonicalIter::new(self, chain, chain_tip, mods)
     }
 
     /// Get a filtered list of outputs from the given `outpoints` that are in `chain` with
@@ -1002,9 +1026,10 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
         outpoints: impl IntoIterator<Item = (OI, OutPoint)> + 'a,
     ) -> impl Iterator<Item = (OI, FullTxOut<A>)> + 'a {
-        self.try_filter_chain_txouts(chain, chain_tip, outpoints)
+        self.try_filter_chain_txouts(chain, chain_tip, mods, outpoints)
             .expect("oracle is infallible")
     }
 
@@ -1030,10 +1055,11 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
         outpoints: impl IntoIterator<Item = (OI, OutPoint)> + 'a,
     ) -> Result<impl Iterator<Item = (OI, FullTxOut<A>)> + 'a, C::Error> {
         Ok(self
-            .try_filter_chain_txouts(chain, chain_tip, outpoints)?
+            .try_filter_chain_txouts(chain, chain_tip, mods, outpoints)?
             .filter(|(_, full_txo)| full_txo.spent_by.is_none()))
     }
 
@@ -1047,9 +1073,10 @@ impl<A: Anchor> TxGraph<A> {
         &'a self,
         chain: &'a C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
         txouts: impl IntoIterator<Item = (OI, OutPoint)> + 'a,
     ) -> impl Iterator<Item = (OI, FullTxOut<A>)> + 'a {
-        self.try_filter_chain_unspents(chain, chain_tip, txouts)
+        self.try_filter_chain_unspents(chain, chain_tip, mods, txouts)
             .expect("oracle is infallible")
     }
 
@@ -1069,6 +1096,7 @@ impl<A: Anchor> TxGraph<A> {
         &self,
         chain: &C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
         outpoints: impl IntoIterator<Item = (OI, OutPoint)>,
         mut trust_predicate: impl FnMut(&OI, ScriptBuf) -> bool,
     ) -> Result<Balance, C::Error> {
@@ -1077,7 +1105,7 @@ impl<A: Anchor> TxGraph<A> {
         let mut untrusted_pending = Amount::ZERO;
         let mut confirmed = Amount::ZERO;
 
-        for (spk_i, txout) in self.try_filter_chain_unspents(chain, chain_tip, outpoints)? {
+        for (spk_i, txout) in self.try_filter_chain_unspents(chain, chain_tip, mods, outpoints)? {
             match &txout.chain_position {
                 ChainPosition::Confirmed { .. } => {
                     if txout.is_confirmed_and_spendable(chain_tip.height) {
@@ -1113,10 +1141,11 @@ impl<A: Anchor> TxGraph<A> {
         &self,
         chain: &C,
         chain_tip: BlockId,
+        mods: CanonicalizationMods,
         outpoints: impl IntoIterator<Item = (OI, OutPoint)>,
         trust_predicate: impl FnMut(&OI, ScriptBuf) -> bool,
     ) -> Balance {
-        self.try_balance(chain, chain_tip, outpoints, trust_predicate)
+        self.try_balance(chain, chain_tip, mods, outpoints, trust_predicate)
             .expect("oracle is infallible")
     }
 }