perf(chain): refactor TxGraph.spends from BTreeMap to HashMap

crates/chain/src/tx_graph.rs

@@ -128,12 +128,12 @@
 use alloc::sync::Arc;
 use alloc::vec::Vec;
 use bdk_core::ConfirmationBlockTime;
 pub use bdk_core::TxUpdate;
 use bitcoin::{Amount, OutPoint, SignedAmount, Transaction, TxOut, Txid};
 use core::fmt::{self, Formatter};
 use core::{
     convert::Infallible,
-    ops::{Deref, RangeInclusive},
+    ops::Deref,
 };
 
 impl<A: Ord> From<TxGraph<A>> for TxUpdate<A> {
@@ -171,7 +171,8 @@
 #[derive(Clone, Debug, PartialEq)]
 pub struct TxGraph<A = ConfirmationBlockTime> {
     txs: HashMap<Txid, TxNodeInternal>,
-    spends: BTreeMap<OutPoint, HashSet<Txid>>,
+    spends: HashMap<OutPoint, HashSet<Txid>>,
+    spend_vouts_by_txid: HashMap<Txid, BTreeSet<u32>>,
     anchors: HashMap<Txid, BTreeSet<A>>,
     first_seen: HashMap<Txid, u64>,
     last_seen: HashMap<Txid, u64>,
@@ -191,6 +192,7 @@
         Self {
             txs: Default::default(),
             spends: Default::default(),
+            spend_vouts_by_txid: Default::default(),
             anchors: Default::default(),
             first_seen: Default::default(),
             last_seen: Default::default(),
@@ -468,11 +470,16 @@
         &self,
         txid: Txid,
     ) -> impl DoubleEndedIterator<Item = (u32, &HashSet<Txid>)> + '_ {
-        let start = OutPoint::new(txid, 0);
-        let end = OutPoint::new(txid, u32::MAX);
-        self.spends
-            .range(start..=end)
-            .map(|(outpoint, spends)| (outpoint.vout, spends))
+        self.spend_vouts_by_txid
+            .get(&txid)
+            .into_iter()
+            .flat_map(move |vouts| {
+                vouts.iter().filter_map(move |vout| {
+                    self.spends
+                        .get(&OutPoint::new(txid, *vout))
+                        .map(|spends| (*vout, spends))
+                })
+            })
     }
 }
 
@@ -707,6 +714,10 @@
                     if txin.previous_output.is_null() {
                         continue;
                     }
+                    self.spend_vouts_by_txid
+                        .entry(txin.previous_output.txid)
+                        .or_default()
+                        .insert(txin.previous_output.vout);
                     self.spends
                         .entry(txin.previous_output)
                         .or_default()
@@ -1417,8 +1428,7 @@
     fn populate_queue(&mut self, depth: usize, txid: Txid) {
         let spend_paths = self
             .graph
-            .spends
-            .range(tx_outpoint_range(txid))
+            .tx_spends(txid)
             .flat_map(|(_, spends)| spends)
             .map(|&txid| (depth, txid));
         self.queue.extend(spend_paths);
@@ -1448,7 +1458,3 @@
         Some(item)
     }
 }
-
-fn tx_outpoint_range(txid: Txid) -> RangeInclusive<OutPoint> {
-    OutPoint::new(txid, u32::MIN)..=OutPoint::new(txid, u32::MAX)
-}

crates/chain/tests/test_tx_graph.rs

@@ -1537,3 +1537,93 @@ fn test_get_first_seen_of_a_tx() {
     let first_seen = graph.get_tx_node(txid).unwrap().first_seen;
     assert_eq!(first_seen, Some(seen_at));
 }
+
+#[test]
+fn test_hashmap_spends_deterministic_vout_order() {
+    let parent_tx = Transaction {
+        version: transaction::Version::ONE,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![],
+        output: vec![TxOut::NULL; 5],
+    };
+    let parent_txid = parent_tx.compute_txid();
+
+    let children: Vec<Transaction> = vec![4u32, 1, 3, 0, 2]
+        .into_iter()
+        .map(|vout| Transaction {
+            version: transaction::Version::ONE,
+            lock_time: absolute::LockTime::ZERO,
+            input: vec![TxIn {
+                previous_output: OutPoint::new(parent_txid, vout),
+                ..Default::default()
+            }],
+            output: vec![],
+        })
+        .collect();
+
+    let mut graph = TxGraph::<ConfirmationBlockTime>::default();
+    let _ = graph.insert_tx(Arc::new(parent_tx));
+    for child in &children {
+        let _ = graph.insert_tx(Arc::new(child.clone()));
+    }
+
+    let vouts_in_order: Vec<u32> = graph.tx_spends(parent_txid).map(|(vout, _)| vout).collect();
+
+    assert_eq!(
+        vouts_in_order,
+        vec![0, 1, 2, 3, 4],
+        "vouts must be returned in sorted order"
+    );
+}
+
+#[test]
+fn test_hashmap_spends_descendant_walking() {
+    let parent_tx = Transaction {
+        version: transaction::Version::ONE,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![],
+        output: vec![TxOut::NULL],
+    };
+    let parent_txid = parent_tx.compute_txid();
+
+    let child = Transaction {
+        version: transaction::Version::ONE,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: OutPoint::new(parent_txid, 0),
+            ..Default::default()
+        }],
+        output: vec![TxOut::NULL],
+    };
+    let child_txid = child.compute_txid();
+
+    let grandchild = Transaction {
+        version: transaction::Version::ONE,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: OutPoint::new(child_txid, 0),
+            ..Default::default()
+        }],
+        output: vec![],
+    };
+    let grandchild_txid = grandchild.compute_txid();
+
+    let mut graph = TxGraph::<ConfirmationBlockTime>::default();
+    let _ = graph.insert_tx(Arc::new(parent_tx));
+    let _ = graph.insert_tx(Arc::new(child));
+    let _ = graph.insert_tx(Arc::new(grandchild));
+
+    let descendants: Vec<Txid> = graph
+        .walk_descendants(parent_txid, |_, txid| Some(txid))
+        .collect();
+
+    assert_eq!(descendants.len(), 2, "parent has 2 descendants");
+    assert!(
+        descendants.contains(&child_txid),
+        "child must be in descendants"
+    );
+    assert!(
+        descendants.contains(&grandchild_txid),
+        "grandchild must be in descendants"
+    );
+}