feat(chain)!: KeychainTxOutIndex uses a universal lookahead

Instead of a per-keychain lookahead, we have a universal lookahead for
all keychains of a `KeychainTxOutIndex` instance. We restrict where the
lookahead can be set to the constructor. This simplifies the API.

Author: evanlinjin

crates/chain/src/keychain/txout_index.rs

@@ -5,7 +5,6 @@ use crate::{
     spk_iter::BIP32_MAX_INDEX,
     SpkIterator, SpkTxOutIndex,
 };
-use alloc::vec::Vec;
 use bitcoin::{OutPoint, Script, TxOut};
 use core::{fmt::Debug, ops::Deref};
 
@@ -67,17 +66,12 @@ pub struct KeychainTxOutIndex<K> {
     // last revealed indexes
     last_revealed: BTreeMap<K, u32>,
     // lookahead settings for each keychain
-    lookahead: BTreeMap<K, u32>,
+    lookahead: u32,
 }
 
 impl<K> Default for KeychainTxOutIndex<K> {
     fn default() -> Self {
-        Self {
-            inner: SpkTxOutIndex::default(),
-            keychains: BTreeMap::default(),
-            last_revealed: BTreeMap::default(),
-            lookahead: BTreeMap::default(),
-        }
+        Self::new(DEFAULT_LOOKAHEAD)
     }
 }
 
@@ -120,6 +114,24 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
     }
 }
 
+impl<K> KeychainTxOutIndex<K> {
+    /// Construct a [`KeychainTxOutIndex`] with the given `lookahead`.
+    ///
+    /// The lookahead is the number of scripts to cache ahead of the last revealed script index.
+    /// This is useful to find outputs you own when processing block data that lie beyond the last
+    /// revealed index. In certain situations, such as when performing an initial scan of the
+    /// blockchain during wallet import, it may be uncertain or unknown what the last revealed index
+    /// is.
+    pub fn new(lookahead: u32) -> Self {
+        Self {
+            inner: SpkTxOutIndex::default(),
+            keychains: BTreeMap::new(),
+            last_revealed: BTreeMap::new(),
+            lookahead,
+        }
+    }
+}
+
 impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// Return a reference to the internal [`SpkTxOutIndex`].
     pub fn inner(&self) -> &SpkTxOutIndex<(K, u32)> {
@@ -136,38 +148,17 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         &self.keychains
     }
 
-    /// Add a keychain to the tracker's `txout_index` with a descriptor to derive addresses and a
-    /// default lookahead of `1_000`.
-    ///
-    /// See [`add_keychain_with_lookahead`] for details.
-    ///
-    /// # Panics
-    ///
-    /// This will panic if a different `descriptor` is introduced to the same `keychain`.
-    ///
-    /// [`add_keychain_with_lookahead`]: Self::add_keychain_with_lookahead
-    pub fn add_keychain(&mut self, keychain: K, descriptor: Descriptor<DescriptorPublicKey>) {
-        self.add_keychain_with_lookahead(keychain, descriptor, DEFAULT_LOOKAHEAD)
-    }
-
     /// Add a keychain to the tracker's `txout_index` with a descriptor to derive addresses.
     ///
     /// Adding a keychain means you will be able to derive new script pubkeys under that keychain
     /// and the txout index will discover transaction outputs with those script pubkeys.
     ///
-    /// Refer to [`set_lookahead`] for an explanation of the `lookahead` parameter.
-    ///
     /// # Panics
     ///
     /// This will panic if a different `descriptor` is introduced to the same `keychain`.
     ///
-    /// [`set_lookahead`]: Self::set_lookahead
-    pub fn add_keychain_with_lookahead(
-        &mut self,
-        keychain: K,
-        descriptor: Descriptor<DescriptorPublicKey>,
-        lookahead: u32,
-    ) {
+    /// [`add_keychain_with_lookahead`]: Self::add_keychain_with_lookahead
+    pub fn add_keychain(&mut self, keychain: K, descriptor: Descriptor<DescriptorPublicKey>) {
         let old_descriptor = &*self
             .keychains
             .entry(keychain.clone())
@@ -176,54 +167,12 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             &descriptor, old_descriptor,
             "keychain already contains a different descriptor"
         );
-        self.set_lookahead(&keychain, lookahead)
-    }
-
-    /// Return the lookahead setting for each keychain.
-    ///
-    /// Refer to [`set_lookahead`] for a deeper explanation of the `lookahead`.
-    ///
-    /// [`set_lookahead`]: Self::set_lookahead
-    pub fn lookaheads(&self) -> &BTreeMap<K, u32> {
-        &self.lookahead
-    }
-
-    /// Convenience method to call [`set_lookahead`] for all keychains.
-    ///
-    /// [`set_lookahead`]: Self::set_lookahead
-    pub fn set_lookahead_for_all(&mut self, lookahead: u32) {
-        for keychain in &self.keychains.keys().cloned().collect::<Vec<_>>() {
-            self.set_lookahead(keychain, lookahead);
-        }
-    }
-
-    /// Set the lookahead count for `keychain`.
-    ///
-    /// The lookahead is the number of scripts to cache ahead of the last revealed script index. This
-    /// is useful to find outputs you own when processing block data that lie beyond the last revealed
-    /// index. In certain situations, such as when performing an initial scan of the blockchain during
-    /// wallet import, it may be uncertain or unknown what the last revealed index is.
-    ///
-    /// # Panics
-    ///
-    /// This will panic if the new `lookahead` value is smaller than the previous value.
-    /// This will also panic if the `keychain` does not exist.
-    pub fn set_lookahead(&mut self, keychain: &K, lookahead: u32) {
-        let old_lookahead = self
-            .lookahead
-            .insert(keychain.clone(), lookahead)
-            .unwrap_or(0);
-        assert!(old_lookahead <= lookahead, "lookahead must always increase");
-        self.replenish_lookahead(keychain);
+        self.replenish_lookahead(&keychain, self.lookahead);
     }
 
-    /// Convenience method to call [`lookahead_to_target`] for multiple keychains.
-    ///
-    /// [`lookahead_to_target`]: Self::lookahead_to_target
-    pub fn lookahead_to_target_multi(&mut self, target_indexes: BTreeMap<K, u32>) {
-        for (keychain, target_index) in target_indexes {
-            self.lookahead_to_target(&keychain, target_index)
-        }
+    /// Return the lookahead setting.
+    pub fn lookahead(&self) -> u32 {
+        self.lookahead
     }
 
     /// Store lookahead scripts until `target_index`.
@@ -232,23 +181,14 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     pub fn lookahead_to_target(&mut self, keychain: &K, target_index: u32) {
         let next_index = self.next_store_index(keychain);
         if let Some(temp_lookahead) = target_index.checked_sub(next_index).filter(|&v| v > 0) {
-            // We temporarily change the lookahead settings (so we can reuse the
-            // `replenish_lookahead` logic).
-            let old_lookahead = self.lookahead.insert(keychain.clone(), temp_lookahead);
-            self.replenish_lookahead(keychain);
-            // revert
-            match old_lookahead {
-                Some(lookahead) => self.lookahead.insert(keychain.clone(), lookahead),
-                None => self.lookahead.remove(keychain),
-            };
+            self.replenish_lookahead(keychain, temp_lookahead);
         }
     }
 
-    fn replenish_lookahead(&mut self, keychain: &K) {
+    fn replenish_lookahead(&mut self, keychain: &K, lookahead: u32) {
         let descriptor = self.keychains.get(keychain).expect("keychain must exist");
         let next_store_index = self.next_store_index(keychain);
         let next_reveal_index = self.last_revealed.get(keychain).map_or(0, |v| *v + 1);
-        let lookahead = self.lookahead.get(keychain).map_or(0, |v| *v);
 
         for (new_index, new_spk) in
             SpkIterator::new_with_range(descriptor, next_store_index..next_reveal_index + lookahead)
@@ -420,22 +360,21 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
 
         let target_index = if has_wildcard { target_index } else { 0 };
         let next_reveal_index = self.last_revealed.get(keychain).map_or(0, |v| *v + 1);
-        let lookahead = self.lookahead.get(keychain).map_or(0, |v| *v);
 
-        debug_assert!(next_reveal_index + lookahead >= self.next_store_index(keychain));
+        debug_assert!(next_reveal_index + self.lookahead >= self.next_store_index(keychain));
 
         // if we need to reveal new indices, the latest revealed index goes here
         let mut reveal_to_index = None;
 
         // if the target is not yet revealed, but is already stored (due to lookahead), we need to
         // set the `reveal_to_index` as target here (as the `for` loop below only updates
         // `reveal_to_index` for indexes that are NOT stored)
-        if next_reveal_index <= target_index && target_index < next_reveal_index + lookahead {
+        if next_reveal_index <= target_index && target_index < next_reveal_index + self.lookahead {
             reveal_to_index = Some(target_index);
         }
 
         // we range over indexes that are not stored
-        let range = next_reveal_index + lookahead..=target_index + lookahead;
+        let range = next_reveal_index + self.lookahead..=target_index + self.lookahead;
         for (new_index, new_spk) in SpkIterator::new_with_range(descriptor, range) {
             let _inserted = self
                 .inner

crates/chain/src/spk_iter.rs

@@ -148,22 +148,14 @@ mod test {
         Descriptor<DescriptorPublicKey>,
         Descriptor<DescriptorPublicKey>,
     ) {
-        let mut txout_index = KeychainTxOutIndex::<TestKeychain>::default();
+        let mut txout_index = KeychainTxOutIndex::<TestKeychain>::new(0);
 
         let secp = Secp256k1::signing_only();
         let (external_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap();
         let (internal_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/*)").unwrap();
 
-        txout_index.add_keychain_with_lookahead(
-            TestKeychain::External,
-            external_descriptor.clone(),
-            0,
-        );
-        txout_index.add_keychain_with_lookahead(
-            TestKeychain::Internal,
-            internal_descriptor.clone(),
-            0,
-        );
+        txout_index.add_keychain(TestKeychain::External, external_descriptor.clone());
+        txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone());
 
         (txout_index, external_descriptor, internal_descriptor)
     }

crates/chain/tests/test_indexed_tx_graph.rs

@@ -27,8 +27,10 @@ fn insert_relevant_txs() {
     let spk_0 = descriptor.at_derivation_index(0).unwrap().script_pubkey();
     let spk_1 = descriptor.at_derivation_index(9).unwrap().script_pubkey();
 
-    let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::default();
-    graph.index.add_keychain_with_lookahead((), descriptor, 10);
+    let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::new(
+        KeychainTxOutIndex::new(10),
+    );
+    graph.index.add_keychain((), descriptor);
 
     let tx_a = Transaction {
         output: vec![
@@ -117,15 +119,12 @@ fn test_list_owned_txouts() {
     let (desc_1, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/0/*)").unwrap();
     let (desc_2, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/1/*)").unwrap();
 
-    let mut graph =
-        IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>::default();
+    let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>::new(
+        KeychainTxOutIndex::new(10),
+    );
 
-    graph
-        .index
-        .add_keychain_with_lookahead("keychain_1".into(), desc_1, 10);
-    graph
-        .index
-        .add_keychain_with_lookahead("keychain_2".into(), desc_2, 10);
+    graph.index.add_keychain("keychain_1".into(), desc_1);
+    graph.index.add_keychain("keychain_2".into(), desc_2);
 
     // Get trusted and untrusted addresses
 

crates/chain/tests/test_keychain_txout_index.rs

@@ -19,29 +19,20 @@ enum TestKeychain {
 }
 
 fn init_txout_index(
-    internal_lookahead: u32,
-    external_lookahead: u32,
+    lookahead: u32,
 ) -> (
     bdk_chain::keychain::KeychainTxOutIndex<TestKeychain>,
     Descriptor<DescriptorPublicKey>,
     Descriptor<DescriptorPublicKey>,
 ) {
-    let mut txout_index = bdk_chain::keychain::KeychainTxOutIndex::<TestKeychain>::default();
+    let mut txout_index = bdk_chain::keychain::KeychainTxOutIndex::<TestKeychain>::new(lookahead);
 
     let secp = bdk_chain::bitcoin::secp256k1::Secp256k1::signing_only();
     let (external_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap();
     let (internal_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/*)").unwrap();
 
-    txout_index.add_keychain_with_lookahead(
-        TestKeychain::External,
-        external_descriptor.clone(),
-        internal_lookahead,
-    );
-    txout_index.add_keychain_with_lookahead(
-        TestKeychain::Internal,
-        internal_descriptor.clone(),
-        external_lookahead,
-    );
+    txout_index.add_keychain(TestKeychain::External, external_descriptor.clone());
+    txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone());
 
     (txout_index, external_descriptor, internal_descriptor)
 }
@@ -57,7 +48,7 @@ fn spk_at_index(descriptor: &Descriptor<DescriptorPublicKey>, index: u32) -> Scr
 fn test_set_all_derivation_indices() {
     use bdk_chain::indexed_tx_graph::Indexer;
 
-    let (mut txout_index, _, _) = init_txout_index(0, 0);
+    let (mut txout_index, _, _) = init_txout_index(0);
     let derive_to: BTreeMap<_, _> =
         [(TestKeychain::External, 12), (TestKeychain::Internal, 24)].into();
     assert_eq!(
@@ -75,11 +66,10 @@ fn test_set_all_derivation_indices() {
 
 #[test]
 fn test_lookahead() {
-    let (mut txout_index, external_desc, internal_desc) = init_txout_index(10, 20);
+    let (mut txout_index, external_desc, internal_desc) = init_txout_index(10);
 
     // given:
     // - external lookahead set to 10
-    // - internal lookahead set to 20
     // when:
     // - set external derivation index to value higher than last, but within the lookahead value
     // expect:
@@ -100,7 +90,7 @@ fn test_lookahead() {
         assert_eq!(
             txout_index.inner().all_spks().len(),
             10 /* external lookahead */ +
-            20 /* internal lookahead */ +
+            10 /* internal lookahead */ +
             index as usize + 1 /* `derived` count */
         );
         assert_eq!(
@@ -130,7 +120,7 @@ fn test_lookahead() {
     }
 
     // given:
-    // - internal lookahead is 20
+    // - internal lookahead is 10
     // - internal derivation index is `None`
     // when:
     // - derivation index is set ahead of current derivation index + lookahead
@@ -151,7 +141,7 @@ fn test_lookahead() {
     assert_eq!(
         txout_index.inner().all_spks().len(),
         10 /* external lookahead */ +
-        20 /* internal lookahead */ +
+        10 /* internal lookahead */ +
         20 /* external stored index count */ +
         25 /* internal stored index count */
     );
@@ -229,7 +219,7 @@ fn test_lookahead() {
 // - last used index should change as expected
 #[test]
 fn test_scan_with_lookahead() {
-    let (mut txout_index, external_desc, _) = init_txout_index(10, 10);
+    let (mut txout_index, external_desc, _) = init_txout_index(10);
 
     let spks: BTreeMap<u32, ScriptBuf> = [0, 10, 20, 30]
         .into_iter()
@@ -283,7 +273,7 @@ fn test_scan_with_lookahead() {
 #[test]
 #[rustfmt::skip]
 fn test_wildcard_derivations() {
-    let (mut txout_index, external_desc, _) = init_txout_index(0, 0);
+    let (mut txout_index, external_desc, _) = init_txout_index(0);
     let external_spk_0 = external_desc.at_derivation_index(0).unwrap().script_pubkey();
     let external_spk_16 = external_desc.at_derivation_index(16).unwrap().script_pubkey();
     let external_spk_26 = external_desc.at_derivation_index(26).unwrap().script_pubkey();
@@ -341,7 +331,7 @@ fn test_wildcard_derivations() {
 
 #[test]
 fn test_non_wildcard_derivations() {
-    let mut txout_index = KeychainTxOutIndex::<TestKeychain>::default();
+    let mut txout_index = KeychainTxOutIndex::<TestKeychain>::new(0);
 
     let secp = bitcoin::secp256k1::Secp256k1::signing_only();
     let (no_wildcard_descriptor, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "wpkh([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/0)").unwrap();
@@ -350,7 +340,7 @@ fn test_non_wildcard_derivations() {
         .unwrap()
         .script_pubkey();
 
-    txout_index.add_keychain_with_lookahead(TestKeychain::External, no_wildcard_descriptor, 0);
+    txout_index.add_keychain(TestKeychain::External, no_wildcard_descriptor);
 
     // given:
     // - `txout_index` with no stored scripts