Merge #1324: [chain] Make KeychainTxOutIndex more range based

fac2283 feat(chain)!: make `KeychainTxOutIndex` more range based (LLFourn)

Pull request description:

  KeychainTxOut index should try and avoid "all" kind of queries. There may be subranges of interest. If the user wants "all" they can just query "..".

  The ideas is that KeychainTxOutIndex should be designed to be able to incorporate many unrelated keychains that can be managed in the same index. We should be able to see the "net_value" of a transaction to a specific subrange. e.g. imagine a collaborative custody service that manages all their user descriptors inside the same `KeychainTxOutIndex`. One user in their service may pay another so when you are analyzing how much a transaction is spending for a particular user you need to do analyze a particular sub-range.

  ### Notes to the reviewers

  - I didn't change `unused_spks` to follow this rule because I want to delete that method some time in the future. `unused_spks` is being used in the examples for syncing but it shouldn't be (the discussion as to why will probably surface in #1194).
  - I haven't applied this reasoning to the methods that return `BTreeMap`s e.g. `all_unbounded_spk_iters`. It probably should be but I haven't made up my mind yet.

  This probably belongs after #1194

  ### Changelog notice

  - `KeychainTxOutIndex` methods modified to take ranges of keychains instead.

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo fmt` and `cargo clippy` before committing

  #### New Features:

  * [x] I've added tests for the new feature
  * [x] I've added docs for the new feature

ACKs for top commit:
  evanlinjin:
    ACK fac2283

Tree-SHA512: ec1e75f19d79f71de4b6d7748ef6da076ca92c2f3fd07e0f0dc88e091bf80c61268880ef78be4bed5e0dbab2572e22028f868f33e68a67d47813195d38d78ba5

Author: evanlinjin

crates/bdk/src/wallet/mod.rs

@@ -1015,7 +1015,7 @@ impl<D> Wallet<D> {
     /// let (sent, received) = wallet.sent_and_received(tx);
     /// ```
     pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) {
-        self.indexed_graph.index.sent_and_received(tx)
+        self.indexed_graph.index.sent_and_received(tx, ..)
     }
 
     /// Get a single transaction from the wallet as a [`CanonicalTx`] (if the transaction exists).

crates/chain/src/keychain/txout_index.rs

@@ -268,15 +268,14 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         self.inner.unmark_used(&(keychain, index))
     }
 
-    /// Computes total input value going from script pubkeys in the index (sent) and the total output
-    /// value going to script pubkeys in the index (received) in `tx`. For the `sent` to be computed
-    /// correctly, the output being spent must have already been scanned by the index. Calculating
-    /// received just uses the [`Transaction`] outputs directly, so it will be correct even if it has
-    /// not been scanned.
-    ///
-    /// This calls [`SpkTxOutIndex::sent_and_received`] internally.
-    pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) {
-        self.inner.sent_and_received(tx)
+    /// Computes the total value transfer effect `tx` has on the script pubkeys belonging to the
+    /// keychains in `range`. Value is *sent* when a script pubkey in the `range` is on an input and
+    /// *received* when it is on an output. For `sent` to be computed correctly, the output being
+    /// spent must have already been scanned by the index. Calculating received just uses the
+    /// [`Transaction`] outputs directly, so it will be correct even if it has not been scanned.
+    pub fn sent_and_received(&self, tx: &Transaction, range: impl RangeBounds<K>) -> (u64, u64) {
+        self.inner
+            .sent_and_received(tx, Self::map_to_inner_bounds(range))
     }
 
     /// Computes the net value that this transaction gives to the script pubkeys in the index and
@@ -286,8 +285,8 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// This calls [`SpkTxOutIndex::net_value`] internally.
     ///
     /// [`sent_and_received`]: Self::sent_and_received
-    pub fn net_value(&self, tx: &Transaction) -> i64 {
-        self.inner.net_value(tx)
+    pub fn net_value(&self, tx: &Transaction, range: impl RangeBounds<K>) -> i64 {
+        self.inner.net_value(tx, Self::map_to_inner_bounds(range))
     }
 }
 
@@ -390,24 +389,32 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             .collect()
     }
 
-    /// Iterate over revealed spks of all keychains.
-    pub fn revealed_spks(&self) -> impl DoubleEndedIterator<Item = (K, u32, &Script)> + Clone {
-        self.keychains.keys().flat_map(|keychain| {
-            self.revealed_keychain_spks(keychain)
-                .map(|(i, spk)| (keychain.clone(), i, spk))
+    /// Iterate over revealed spks of keychains in `range`
+    pub fn revealed_spks(
+        &self,
+        range: impl RangeBounds<K>,
+    ) -> impl DoubleEndedIterator<Item = (&K, u32, &Script)> + Clone {
+        self.keychains.range(range).flat_map(|(keychain, _)| {
+            let start = Bound::Included((keychain.clone(), u32::MIN));
+            let end = match self.last_revealed.get(keychain) {
+                Some(last_revealed) => Bound::Included((keychain.clone(), *last_revealed)),
+                None => Bound::Excluded((keychain.clone(), u32::MIN)),
+            };
+
+            self.inner
+                .all_spks()
+                .range((start, end))
+                .map(|((keychain, i), spk)| (keychain, *i, spk.as_script()))
         })
     }
 
     /// Iterate over revealed spks of the given `keychain`.
-    pub fn revealed_keychain_spks(
-        &self,
-        keychain: &K,
-    ) -> impl DoubleEndedIterator<Item = (u32, &Script)> + Clone {
-        let next_i = self.last_revealed.get(keychain).map_or(0, |&i| i + 1);
-        self.inner
-            .all_spks()
-            .range((keychain.clone(), u32::MIN)..(keychain.clone(), next_i))
-            .map(|((_, i), spk)| (*i, spk.as_script()))
+    pub fn revealed_keychain_spks<'a>(
+        &'a self,
+        keychain: &'a K,
+    ) -> impl DoubleEndedIterator<Item = (u32, &Script)> + 'a {
+        self.revealed_spks(keychain..=keychain)
+            .map(|(_, i, spk)| (i, spk))
     }
 
     /// Iterate over revealed, but unused, spks of all keychains.
@@ -617,38 +624,40 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         }
     }
 
-    /// Iterate over all [`OutPoint`]s that point to `TxOut`s with script pubkeys derived from
+    /// Iterate over all [`OutPoint`]s that have `TxOut`s with script pubkeys derived from
     /// `keychain`.
-    ///
-    /// Use [`keychain_outpoints_in_range`](KeychainTxOutIndex::keychain_outpoints_in_range) to
-    /// iterate over a specific derivation range.
-    pub fn keychain_outpoints(
-        &self,
-        keychain: &K,
-    ) -> impl DoubleEndedIterator<Item = (u32, OutPoint)> + '_ {
-        self.keychain_outpoints_in_range(keychain, ..)
+    pub fn keychain_outpoints<'a>(
+        &'a self,
+        keychain: &'a K,
+    ) -> impl DoubleEndedIterator<Item = (u32, OutPoint)> + 'a {
+        self.keychain_outpoints_in_range(keychain..=keychain)
+            .map(move |(_, i, op)| (i, op))
+    }
+
+    /// Iterate over [`OutPoint`]s that have script pubkeys derived from keychains in `range`.
+    pub fn keychain_outpoints_in_range<'a>(
+        &'a self,
+        range: impl RangeBounds<K> + 'a,
+    ) -> impl DoubleEndedIterator<Item = (&'a K, u32, OutPoint)> + 'a {
+        let bounds = Self::map_to_inner_bounds(range);
+        self.inner
+            .outputs_in_range(bounds)
+            .map(move |((keychain, i), op)| (keychain, *i, op))
     }
 
-    /// Iterate over [`OutPoint`]s that point to `TxOut`s with script pubkeys derived from
-    /// `keychain` in a given derivation `range`.
-    pub fn keychain_outpoints_in_range(
-        &self,
-        keychain: &K,
-        range: impl RangeBounds<u32>,
-    ) -> impl DoubleEndedIterator<Item = (u32, OutPoint)> + '_ {
-        let start = match range.start_bound() {
-            Bound::Included(i) => Bound::Included((keychain.clone(), *i)),
-            Bound::Excluded(i) => Bound::Excluded((keychain.clone(), *i)),
+    fn map_to_inner_bounds(bound: impl RangeBounds<K>) -> impl RangeBounds<(K, u32)> {
+        let start = match bound.start_bound() {
+            Bound::Included(keychain) => Bound::Included((keychain.clone(), u32::MIN)),
+            Bound::Excluded(keychain) => Bound::Excluded((keychain.clone(), u32::MAX)),
             Bound::Unbounded => Bound::Unbounded,
         };
-        let end = match range.end_bound() {
-            Bound::Included(i) => Bound::Included((keychain.clone(), *i)),
-            Bound::Excluded(i) => Bound::Excluded((keychain.clone(), *i)),
+        let end = match bound.end_bound() {
+            Bound::Included(keychain) => Bound::Included((keychain.clone(), u32::MAX)),
+            Bound::Excluded(keychain) => Bound::Excluded((keychain.clone(), u32::MIN)),
             Bound::Unbounded => Bound::Unbounded,
         };
-        self.inner
-            .outputs_in_range((start, end))
-            .map(|((_, i), op)| (*i, op))
+
+        (start, end)
     }
 
     /// Returns the highest derivation index of the `keychain` where [`KeychainTxOutIndex`] has

crates/chain/src/spk_txout_index.rs

@@ -270,36 +270,39 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
         self.spk_indices.get(script)
     }
 
-    /// Computes total input value going from script pubkeys in the index (sent) and the total output
-    /// value going to script pubkeys in the index (received) in `tx`. For the `sent` to be computed
-    /// correctly, the output being spent must have already been scanned by the index. Calculating
-    /// received just uses the [`Transaction`] outputs directly, so it will be correct even if it has
-    /// not been scanned.
-    pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) {
+    /// Computes the total value transfer effect `tx` has on the script pubkeys in `range`. Value is
+    /// *sent* when a script pubkey in the `range` is on an input and *received* when it is on an
+    /// output. For `sent` to be computed correctly, the output being spent must have already been
+    /// scanned by the index. Calculating received just uses the [`Transaction`] outputs directly,
+    /// so it will be correct even if it has not been scanned.
+    pub fn sent_and_received(&self, tx: &Transaction, range: impl RangeBounds<I>) -> (u64, u64) {
         let mut sent = 0;
         let mut received = 0;
 
         for txin in &tx.input {
-            if let Some((_, txout)) = self.txout(txin.previous_output) {
-                sent += txout.value.to_sat();
+            if let Some((index, txout)) = self.txout(txin.previous_output) {
+                if range.contains(index) {
+                    sent += txout.value.to_sat();
+                }
             }
         }
         for txout in &tx.output {
-            if self.index_of_spk(&txout.script_pubkey).is_some() {
-                received += txout.value.to_sat();
+            if let Some(index) = self.index_of_spk(&txout.script_pubkey) {
+                if range.contains(index) {
+                    received += txout.value.to_sat();
+                }
             }
         }
 
         (sent, received)
     }
 
-    /// Computes the net value that this transaction gives to the script pubkeys in the index and
-    /// *takes* from the transaction outputs in the index. Shorthand for calling
-    /// [`sent_and_received`] and subtracting sent from received.
+    /// Computes the net value transfer effect of `tx` on the script pubkeys in `range`. Shorthand
+    /// for calling [`sent_and_received`] and subtracting sent from received.
     ///
     /// [`sent_and_received`]: Self::sent_and_received
-    pub fn net_value(&self, tx: &Transaction) -> i64 {
-        let (sent, received) = self.sent_and_received(tx);
+    pub fn net_value(&self, tx: &Transaction, range: impl RangeBounds<I>) -> i64 {
+        let (sent, received) = self.sent_and_received(tx, range);
         received as i64 - sent as i64
     }
 

crates/chain/tests/test_spk_txout_index.rs

@@ -20,11 +20,13 @@ fn spk_txout_sent_and_received() {
         }],
     };
 
-    assert_eq!(index.sent_and_received(&tx1), (0, 42_000));
-    assert_eq!(index.net_value(&tx1), 42_000);
+    assert_eq!(index.sent_and_received(&tx1, ..), (0, 42_000));
+    assert_eq!(index.sent_and_received(&tx1, ..1), (0, 42_000));
+    assert_eq!(index.sent_and_received(&tx1, 1..), (0, 0));
+    assert_eq!(index.net_value(&tx1, ..), 42_000);
     index.index_tx(&tx1);
     assert_eq!(
-        index.sent_and_received(&tx1),
+        index.sent_and_received(&tx1, ..),
         (0, 42_000),
         "shouldn't change after scanning"
     );
@@ -51,8 +53,10 @@ fn spk_txout_sent_and_received() {
         ],
     };
 
-    assert_eq!(index.sent_and_received(&tx2), (42_000, 50_000));
-    assert_eq!(index.net_value(&tx2), 8_000);
+    assert_eq!(index.sent_and_received(&tx2, ..), (42_000, 50_000));
+    assert_eq!(index.sent_and_received(&tx2, ..1), (42_000, 30_000));
+    assert_eq!(index.sent_and_received(&tx2, 1..), (0, 20_000));
+    assert_eq!(index.net_value(&tx2, ..), 8_000);
 }
 
 #[test]

example-crates/example_electrum/src/main.rs

@@ -210,8 +210,8 @@ fn main() -> anyhow::Result<()> {
             if all_spks {
                 let all_spks = graph
                     .index
-                    .revealed_spks()
-                    .map(|(k, i, spk)| (k, i, spk.to_owned()))
+                    .revealed_spks(..)
+                    .map(|(k, i, spk)| (k.to_owned(), i, spk.to_owned()))
                     .collect::<Vec<_>>();
                 spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| {
                     eprintln!("scanning {}:{}", k, i);

example-crates/example_esplora/src/main.rs

@@ -241,8 +241,8 @@ fn main() -> anyhow::Result<()> {
                 if *all_spks {
                     let all_spks = graph
                         .index
-                        .revealed_spks()
-                        .map(|(k, i, spk)| (k, i, spk.to_owned()))
+                        .revealed_spks(..)
+                        .map(|(k, i, spk)| (k.to_owned(), i, spk.to_owned()))
                         .collect::<Vec<_>>();
                     spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| {
                         eprintln!("scanning {}:{}", k, i);