fix: TxGraph::missing_blocks logic

evanlinjin · evanlinjin · commit bea8e5aff4f1 · 2023-07-28T11:37:17.000+08:00
Added additional tests for unnecessary duplicate heights
diff --git a/crates/chain/src/tx_data_traits.rs b/crates/chain/src/tx_data_traits.rs
@@ -38,6 +38,8 @@ impl ForEachTxOut for Transaction {
 
 /// Trait that "anchors" blockchain data to a specific block of height and hash.
 ///
+/// [`Anchor`] implementations must be [`Ord`] by the anchor block's [`BlockId`] first.
+///
 /// I.e. If transaction A is anchored in block B, then if block B is in the best chain, we can
 /// assume that transaction A is also confirmed in the best chain. This does not necessarily mean
 /// that transaction A is confirmed in block B. It could also mean transaction A is confirmed in a
diff --git a/crates/chain/src/tx_graph.rs b/crates/chain/src/tx_graph.rs
@@ -601,23 +601,63 @@ impl<A: Anchor> TxGraph<A> {
     /// Find missing block heights of `chain`.
     ///
     /// This works by scanning through anchors, and seeing whether the anchor block of the anchor
-    /// exists in the [`LocalChain`].
-    pub fn missing_blocks<'a>(&'a self, chain: &'a LocalChain) -> impl Iterator<Item = u32> + 'a {
-        let mut last_block = Option::<BlockId>::None;
+    /// exists in the [`LocalChain`]. The returned iterator does not output duplicate heights.
+    pub fn missing_heights<'a>(&'a self, chain: &'a LocalChain) -> impl Iterator<Item = u32> + 'a {
+        // Map of txids to skip.
+        //
+        // Usually, if a height of a tx anchor is missing from the chain, we would want to return
+        // this height in the iterator. The exception is when the tx is confirmed in chain. All the
+        // other missing-height anchors of this tx can be skipped.
+        //
+        // * Some(true)  => skip all anchors of this txid
+        // * Some(false) => do not skip anchors of this txid
+        // * None        => we do not know whether we can skip this txid
+        let mut txids_to_skip = HashMap::<Txid, bool>::new();
+
+        // Keeps track of the last height emitted so we don't double up.
+        let mut last_height_emitted = Option::<u32>::None;
+
         self.anchors
             .iter()
-            .map(|(a, _)| a.anchor_block())
-            .filter(move |block| {
-                if last_block.as_ref() == Some(block) {
-                    false
-                } else {
-                    last_block = Some(*block);
+            .filter(move |(_, txid)| {
+                let skip = *txids_to_skip.entry(*txid).or_insert_with(|| {
+                    let tx_anchors = match self.txs.get(txid) {
+                        Some((_, anchors, _)) => anchors,
+                        None => return true,
+                    };
+                    let mut has_missing_height = false;
+                    for anchor_block in tx_anchors.iter().map(Anchor::anchor_block) {
+                        match chain.heights().get(&anchor_block.height) {
+                            None => {
+                                has_missing_height = true;
+                                continue;
+                            }
+                            Some(chain_hash) => {
+                                if chain_hash == &anchor_block.hash {
+                                    return true;
+                                }
+                            }
+                        }
+                    }
+                    !has_missing_height
+                });
+                #[cfg(feature = "std")]
+                debug_assert!({
+                    println!("txid={} skip={}", txid, skip);
                     true
-                }
+                });
+                !skip
             })
-            .filter_map(|block| match chain.heights().get(&block.height) {
-                Some(chain_hash) if *chain_hash == block.hash => None,
-                _ => Some(block.height),
+            .filter_map(move |(a, _)| {
+                let anchor_block = a.anchor_block();
+                if Some(anchor_block.height) != last_height_emitted
+                    && !chain.heights().contains_key(&anchor_block.height)
+                {
+                    last_height_emitted = Some(anchor_block.height);
+                    Some(anchor_block.height)
+                } else {
+                    None
+                }
             })
     }
 
diff --git a/crates/chain/tests/test_tx_graph.rs b/crates/chain/tests/test_tx_graph.rs
@@ -4,7 +4,7 @@ use bdk_chain::{
     collections::*,
     local_chain::LocalChain,
     tx_graph::{Additions, TxGraph},
-    Append, BlockId, ChainPosition, ConfirmationHeightAnchor,
+    Anchor, Append, BlockId, ChainPosition, ConfirmationHeightAnchor,
 };
 use bitcoin::{
     hashes::Hash, BlockHash, OutPoint, PackedLockTime, Script, Transaction, TxIn, TxOut, Txid,
@@ -822,3 +822,136 @@ fn test_additions_last_seen_append() {
         );
     }
 }
+
+#[test]
+fn test_missing_blocks() {
+    /// An anchor implementation for testing, made up of `(the_anchor_block, random_data)`.
+    #[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, core::hash::Hash)]
+    struct TestAnchor(BlockId);
+
+    impl Anchor for TestAnchor {
+        fn anchor_block(&self) -> BlockId {
+            self.0
+        }
+    }
+
+    struct Scenario<'a> {
+        name: &'a str,
+        graph: TxGraph<TestAnchor>,
+        chain: LocalChain,
+        exp_heights: &'a [u32],
+    }
+
+    const fn new_anchor(height: u32, hash: BlockHash) -> TestAnchor {
+        TestAnchor(BlockId { height, hash })
+    }
+
+    fn new_scenario<'a>(
+        name: &'a str,
+        graph_anchors: &'a [(Txid, TestAnchor)],
+        chain: &'a [(u32, BlockHash)],
+        exp_heights: &'a [u32],
+    ) -> Scenario<'a> {
+        Scenario {
+            name,
+            graph: {
+                let mut g = TxGraph::default();
+                for (txid, anchor) in graph_anchors {
+                    let _ = g.insert_anchor(*txid, anchor.clone());
+                }
+                g
+            },
+            chain: {
+                let mut c = LocalChain::default();
+                for (height, hash) in chain {
+                    let _ = c.insert_block(BlockId {
+                        height: *height,
+                        hash: *hash,
+                    });
+                }
+                c
+            },
+            exp_heights,
+        }
+    }
+
+    fn run(scenarios: &[Scenario]) {
+        for scenario in scenarios {
+            let Scenario {
+                name,
+                graph,
+                chain,
+                exp_heights,
+            } = scenario;
+
+            let heights = graph.missing_heights(chain).collect::<Vec<_>>();
+            assert_eq!(&heights, exp_heights, "scenario: {}", name);
+        }
+    }
+
+    run(&[
+        new_scenario(
+            "2 txs with the same anchor (2:B) which is missing from chain",
+            &[
+                (h!("tx_1"), new_anchor(2, h!("B"))),
+                (h!("tx_2"), new_anchor(2, h!("B"))),
+            ],
+            &[(1, h!("A")), (3, h!("C"))],
+            &[2],
+        ),
+        new_scenario(
+            "2 txs with different anchors at the same height, one of the anchors is missing",
+            &[
+                (h!("tx_1"), new_anchor(2, h!("B1"))),
+                (h!("tx_2"), new_anchor(2, h!("B2"))),
+            ],
+            &[(1, h!("A")), (2, h!("B1"))],
+            &[],
+        ),
+        new_scenario(
+            "tx with 2 anchors of same height which are missing from the chain",
+            &[
+                (h!("tx"), new_anchor(3, h!("C1"))),
+                (h!("tx"), new_anchor(3, h!("C2"))),
+            ],
+            &[(1, h!("A")), (4, h!("D"))],
+            &[3],
+        ),
+        new_scenario(
+            "tx with 2 anchors at the same height, chain has this height but does not match either anchor",
+            &[
+                (h!("tx"), new_anchor(4, h!("D1"))),
+                (h!("tx"), new_anchor(4, h!("D2"))),
+            ],
+            &[(4, h!("D3")), (5, h!("E"))],
+            &[],
+        ),
+        new_scenario(
+            "tx with 2 anchors at different heights, one anchor exists in chain, should return nothing",
+            &[
+                (h!("tx"), new_anchor(3, h!("C"))),
+                (h!("tx"), new_anchor(4, h!("D"))),
+            ],
+            &[(4, h!("D")), (5, h!("E"))],
+            &[],
+        ),
+        new_scenario(
+            "tx with 2 anchors at different heights, first height is already in chain with different hash, iterator should only return 2nd height",
+            &[
+                (h!("tx"), new_anchor(5, h!("E1"))),
+                (h!("tx"), new_anchor(6, h!("F1"))),
+            ],
+            &[(4, h!("D")), (5, h!("E")), (7, h!("G"))],
+            &[6],
+        ),
+        new_scenario(
+            "tx with 2 anchors at different heights, neither height is in chain, both heights should be returned",
+            &[
+                (h!("tx"), new_anchor(3, h!("C"))),
+                (h!("tx"), new_anchor(4, h!("D"))),
+            ],
+            &[(1, h!("A")), (2, h!("B"))],
+            &[3, 4],
+        ),
+    ]);
+}
diff --git a/example-crates/wallet_esplora/src/main.rs b/example-crates/wallet_esplora/src/main.rs
@@ -56,7 +56,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let (update_graph, last_active_indices) =
         client.update_tx_graph(keychain_spks, None, None, STOP_GAP, PARALLEL_REQUESTS)?;
-    let get_heights = wallet.tx_graph().missing_blocks(wallet.local_chain());
+    let get_heights = wallet.tx_graph().missing_heights(wallet.local_chain());
     let chain_update = client.update_local_chain(prev_tip, get_heights)?;
     let update = LocalUpdate {
         last_active_indices,
diff --git a/example-crates/wallet_esplora_async/src/main.rs b/example-crates/wallet_esplora_async/src/main.rs
@@ -57,7 +57,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let (update_graph, last_active_indices) = client
         .update_tx_graph(keychain_spks, None, None, STOP_GAP, PARALLEL_REQUESTS)
         .await?;
-    let get_heights = wallet.tx_graph().missing_blocks(wallet.local_chain());
+    let get_heights = wallet.tx_graph().missing_heights(wallet.local_chain());
     let chain_update = client.update_local_chain(prev_tip, get_heights).await?;
     let update = LocalUpdate {
         last_active_indices,
