feat(chain)!: KeychainTxOutIndex uses a universal lookahead

The wallet is currently created without setting any lookahead value for the keychain. This implicitly makes it a lookahead of 0. As this is a high-level interface we should avoid footguns and aim for a reasonable default. Instead of simply patching it for wallet, we alter `KeychainTxOutIndex` to have a default lookahead value. Additionally, instead of a per-keychain lookahead, the constructor asks for a `lookahead` value. This avoids the footguns of having methods which allows the caller the decrease the `lookahead` (and therefore panicing). This also simplifies the API. Co-authored-by: Antoine Poisot <[email protected]> Co-authored-by: 志宇 <[email protected]>
bitcoindevkit · Dec 21, 2023 · 627953d · 627953d
1 parent b5612f2
commit 627953d
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 103 deletions.
diff --git a/crates/chain/src/keychain/txout_index.rs b/crates/chain/src/keychain/txout_index.rs
@@ -5,12 +5,13 @@ use crate::{
     spk_iter::BIP32_MAX_INDEX,
     SpkIterator, SpkTxOutIndex,
 };
-use alloc::vec::Vec;
 use bitcoin::{OutPoint, Script, TxOut};
 use core::{fmt::Debug, ops::Deref};
 
 use crate::Append;
 
+const DEFAULT_LOOKAHEAD: u32 = 1_000;
+
 /// A convenient wrapper around [`SpkTxOutIndex`] that relates script pubkeys to miniscript public
 /// [`Descriptor`]s.
 ///
@@ -46,7 +47,7 @@ use crate::Append;
 /// # 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();
-/// # let descriptor_for_user_42 = external_descriptor.clone();
+/// # let (descriptor_for_user_42, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/2/*)").unwrap();
 /// txout_index.add_keychain(MyKeychain::External, external_descriptor);
 /// txout_index.add_keychain(MyKeychain::Internal, internal_descriptor);
 /// txout_index.add_keychain(MyKeychain::MyAppUser { user_id: 42 }, descriptor_for_user_42);
@@ -65,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)
     }
 }
 
@@ -118,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)> {
@@ -142,57 +156,27 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// # 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>) {
         let old_descriptor = &*self
             .keychains
-            .entry(keychain)
+            .entry(keychain.clone())
             .or_insert_with(|| descriptor.clone());
         assert_eq!(
             &descriptor, old_descriptor,
             "keychain already contains a different descriptor"
         );
+        self.replenish_lookahead(&keychain, self.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.
+    /// Get the lookahead setting.
     ///
-    /// # Panics
+    /// Refer to [`new`] for more information on the `lookahead`.
     ///
-    /// This will panic if the `keychain` does not exist.
-    pub fn set_lookahead(&mut self, keychain: &K, lookahead: u32) {
-        self.lookahead.insert(keychain.clone(), lookahead);
-        self.replenish_lookahead(keychain);
-    }
-
-    /// 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)
-        }
+    /// [`new`]: Self::new
+    pub fn lookahead(&self) -> u32 {
+        self.lookahead
     }
 
     /// Store lookahead scripts until `target_index`.
@@ -201,22 +185,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) {
-            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)
@@ -388,25 +364,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_eq!(
-            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

diff --git a/crates/chain/src/spk_iter.rs b/crates/chain/src/spk_iter.rs
@@ -148,7 +148,7 @@ 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();

diff --git a/crates/chain/tests/test_indexed_tx_graph.rs b/crates/chain/tests/test_indexed_tx_graph.rs
@@ -27,9 +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();
+    let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::new(
+        KeychainTxOutIndex::new(10),
+    );
     graph.index.add_keychain((), descriptor);
-    graph.index.set_lookahead(&(), 10);
 
     let tx_a = Transaction {
         output: vec![
@@ -118,12 +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("keychain_1".into(), desc_1);
     graph.index.add_keychain("keychain_2".into(), desc_2);
-    graph.index.set_lookahead_for_all(10);
 
     // Get trusted and untrusted addresses
 

diff --git a/crates/chain/tests/test_keychain_txout_index.rs b/crates/chain/tests/test_keychain_txout_index.rs
@@ -18,12 +18,14 @@ enum TestKeychain {
     Internal,
 }
 
-fn init_txout_index() -> (
+fn init_txout_index(
+    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();
@@ -46,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();
+    let (mut txout_index, _, _) = init_txout_index(0);
     let derive_to: BTreeMap<_, _> =
         [(TestKeychain::External, 12), (TestKeychain::Internal, 24)].into();
     assert_eq!(
@@ -64,19 +66,10 @@ fn test_set_all_derivation_indices() {
 
 #[test]
 fn test_lookahead() {
-    let (mut txout_index, external_desc, internal_desc) = init_txout_index();
-
-    // ensure it does not break anything if lookahead is set multiple times
-    (0..=10).for_each(|lookahead| txout_index.set_lookahead(&TestKeychain::External, lookahead));
-    (0..=20)
-        .filter(|v| v % 2 == 0)
-        .for_each(|lookahead| txout_index.set_lookahead(&TestKeychain::Internal, lookahead));
-
-    assert_eq!(txout_index.inner().all_spks().len(), 30);
+    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:
@@ -97,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!(
@@ -127,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
@@ -148,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 */
     );
@@ -226,8 +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();
-    txout_index.set_lookahead_for_all(10);
+    let (mut txout_index, external_desc, _) = init_txout_index(10);
 
     let spks: BTreeMap<u32, ScriptBuf> = [0, 10, 20, 30]
         .into_iter()
@@ -281,7 +273,7 @@ fn test_scan_with_lookahead() {
 #[test]
 #[rustfmt::skip]
 fn test_wildcard_derivations() {
-    let (mut txout_index, external_desc, _) = init_txout_index();
+    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();
@@ -339,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();

diff --git a/example-crates/example_bitcoind_rpc_polling/src/main.rs b/example-crates/example_bitcoind_rpc_polling/src/main.rs
@@ -64,9 +64,6 @@ struct RpcArgs {
     /// Starting block height to fallback to if no point of agreement if found
     #[clap(env = "FALLBACK_HEIGHT", long, default_value = "0")]
     fallback_height: u32,
-    /// The unused-scripts lookahead will be kept at this size
-    #[clap(long, default_value = "10")]
-    lookahead: u32,
 }
 
 impl From<RpcArgs> for Auth {
@@ -161,13 +158,9 @@ fn main() -> anyhow::Result<()> {
     match rpc_cmd {
         RpcCommands::Sync { rpc_args } => {
             let RpcArgs {
-                fallback_height,
-                lookahead,
-                ..
+                fallback_height, ..
             } = rpc_args;
 
-            graph.lock().unwrap().index.set_lookahead_for_all(lookahead);
-
             let chain_tip = chain.lock().unwrap().tip();
             let rpc_client = rpc_args.new_client()?;
             let mut emitter = Emitter::new(&rpc_client, chain_tip, fallback_height);
@@ -233,13 +226,10 @@ fn main() -> anyhow::Result<()> {
         }
         RpcCommands::Live { rpc_args } => {
             let RpcArgs {
-                fallback_height,
-                lookahead,
-                ..
+                fallback_height, ..
             } = rpc_args;
             let sigterm_flag = start_ctrlc_handler();
 
-            graph.lock().unwrap().index.set_lookahead_for_all(lookahead);
             let last_cp = chain.lock().unwrap().tip();
 
             println!(