Implement SpkIterator

SpkIterator was created with its own nth() and next() implementations
and its own new() and new_with_range() constructors.

Co-authored-by: 志宇 <hello@evanlinjin.me>

crates/chain/src/keychain/txout_index.rs

@@ -2,19 +2,17 @@ use crate::{
     collections::*,
     indexed_tx_graph::{Indexer, OwnedIndexer},
     miniscript::{Descriptor, DescriptorPublicKey},
-    ForEachTxOut, SpkTxOutIndex,
+    spk_iter::BIP32_MAX_INDEX,
+    ForEachTxOut, SpkIterator, SpkTxOutIndex,
 };
-use alloc::{borrow::Cow, vec::Vec};
+use alloc::vec::Vec;
-use bitcoin::{secp256k1::Secp256k1, OutPoint, Script, TxOut};
+use bitcoin::{OutPoint, Script, TxOut};
 use core::{fmt::Debug, ops::Deref};
 
 use crate::Append;
 
 use super::DerivationAdditions;
 
-/// Maximum [BIP32](https://bips.xyz/32) derivation index.
-pub const BIP32_MAX_INDEX: u32 = (1 << 31) - 1;
-
 /// A convenient wrapper around [`SpkTxOutIndex`] that relates script pubkeys to miniscript public
 /// [`Descriptor`]s.
 ///
@@ -243,10 +241,9 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         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 range_descriptor_spks(
+        for (new_index, new_spk) in
-            Cow::Borrowed(descriptor),
+            SpkIterator::new_with_range(descriptor, next_store_index..next_reveal_index + lookahead)
-            next_store_index..next_reveal_index + lookahead,
+        {
-        ) {
             let _inserted = self
                 .inner
                 .insert_spk((keychain.clone(), new_index), new_spk);
@@ -266,13 +263,13 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// derivable script pubkeys.
     pub fn spks_of_all_keychains(
         &self,
-    ) -> BTreeMap<K, impl Iterator<Item = (u32, Script)> + Clone> {
+    ) -> BTreeMap<K, SpkIterator<Descriptor<DescriptorPublicKey>>> {
         self.keychains
             .iter()
             .map(|(keychain, descriptor)| {
                 (
                     keychain.clone(),
-                    range_descriptor_spks(Cow::Owned(descriptor.clone()), 0..),
+                    SpkIterator::new_with_range(descriptor.clone(), 0..),
                 )
             })
             .collect()
@@ -284,13 +281,13 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// # Panics
     ///
     /// This will panic if the `keychain` does not exist.
-    pub fn spks_of_keychain(&self, keychain: &K) -> impl Iterator<Item = (u32, Script)> + Clone {
+    pub fn spks_of_keychain(&self, keychain: &K) -> SpkIterator<Descriptor<DescriptorPublicKey>> {
         let descriptor = self
             .keychains
             .get(keychain)
             .expect("keychain must exist")
             .clone();
-        range_descriptor_spks(Cow::Owned(descriptor), 0..)
+        SpkIterator::new_with_range(descriptor, 0..)
     }
 
     /// Convenience method to get [`revealed_spks_of_keychain`] of all keychains.
@@ -370,7 +367,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         &mut self,
         keychains: &BTreeMap<K, u32>,
     ) -> (
-        BTreeMap<K, impl Iterator<Item = (u32, Script)>>,
+        BTreeMap<K, SpkIterator<Descriptor<DescriptorPublicKey>>>,
         DerivationAdditions<K>,
     ) {
         let mut additions = DerivationAdditions::default();
@@ -380,7 +377,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             let (new_spks, new_additions) = self.reveal_to_target(keychain, index);
             if !new_additions.is_empty() {
                 spks.insert(keychain.clone(), new_spks);
-                additions.append(new_additions);
+                additions.append(new_additions.clone());
             }
         }
 
@@ -405,7 +402,10 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         &mut self,
         keychain: &K,
         target_index: u32,
-    ) -> (impl Iterator<Item = (u32, Script)>, DerivationAdditions<K>) {
+    ) -> (
+        SpkIterator<Descriptor<DescriptorPublicKey>>,
+        DerivationAdditions<K>,
+    ) {
         let descriptor = self.keychains.get(keychain).expect("keychain must exist");
         let has_wildcard = descriptor.has_wildcard();
 
@@ -430,7 +430,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
 
         // we range over indexes that are not stored
         let range = next_reveal_index + lookahead..=target_index + lookahead;
-        for (new_index, new_spk) in range_descriptor_spks(Cow::Borrowed(descriptor), range) {
+        for (new_index, new_spk) in SpkIterator::new_with_range(descriptor, range) {
             let _inserted = self
                 .inner
                 .insert_spk((keychain.clone(), new_index), new_spk);
@@ -447,16 +447,13 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
                 let _old_index = self.last_revealed.insert(keychain.clone(), index);
                 debug_assert!(_old_index < Some(index));
                 (
-                    range_descriptor_spks(
+                    SpkIterator::new_with_range(descriptor.clone(), next_reveal_index..index + 1),
-                        Cow::Owned(descriptor.clone()),
-                        next_reveal_index..index + 1,
-                    ),
                     DerivationAdditions(core::iter::once((keychain.clone(), index)).collect()),
                 )
             }
             None => (
-                range_descriptor_spks(
+                SpkIterator::new_with_range(
-                    Cow::Owned(descriptor.clone()),
+                    descriptor.clone(),
                     next_reveal_index..next_reveal_index,
                 ),
                 DerivationAdditions::default(),
@@ -587,33 +584,3 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         let _ = self.reveal_to_target_multi(&additions.0);
     }
 }
-
-fn range_descriptor_spks<'a, R>(
-    descriptor: Cow<'a, Descriptor<DescriptorPublicKey>>,
-    range: R,
-) -> impl Iterator<Item = (u32, Script)> + Clone + Send + 'a
-where
-    R: Iterator<Item = u32> + Clone + Send + 'a,
-{
-    let secp = Secp256k1::verification_only();
-    let has_wildcard = descriptor.has_wildcard();
-    range
-        .into_iter()
-        // non-wildcard descriptors can only have one derivation index (0)
-        .take_while(move |&index| has_wildcard || index == 0)
-        // we can only iterate over non-hardened indices
-        .take_while(|&index| index <= BIP32_MAX_INDEX)
-        .map(
-            move |index| -> Result<_, miniscript::descriptor::ConversionError> {
-                Ok((
-                    index,
-                    descriptor
-                        .at_derivation_index(index)
-                        .derived_descriptor(&secp)?
-                        .script_pubkey(),
-                ))
-            },
-        )
-        .take_while(Result::is_ok)
-        .map(Result::unwrap)
-}

crates/chain/src/lib.rs

@@ -43,6 +43,10 @@ pub use miniscript;
 mod descriptor_ext;
 #[cfg(feature = "miniscript")]
 pub use descriptor_ext::DescriptorExt;
+#[cfg(feature = "miniscript")]
+mod spk_iter;
+#[cfg(feature = "miniscript")]
+pub use spk_iter::*;
 
 #[allow(unused_imports)]
 #[macro_use]

crates/chain/src/spk_iter.rs

@@ -0,0 +1,215 @@
+use crate::{
+    bitcoin::{secp256k1::Secp256k1, Script},
+    miniscript::{Descriptor, DescriptorPublicKey},
+};
+use core::{borrow::Borrow, ops::Bound, ops::RangeBounds};
+
+/// Maximum [BIP32](https://bips.xyz/32) derivation index.
+pub const BIP32_MAX_INDEX: u32 = (1 << 31) - 1;
+
+/// An iterator for derived script pubkeys.
+///
+/// [`SpkIterator`] is an implementation of the [`Iterator`] trait which possesses its own `next()`
+/// and `nth()` functions, both of which circumvent the unnecessary intermediate derivations required
+/// when using their default implementations.
+///
+/// ## Examples
+///
+/// ```
+/// use bdk_chain::SpkIterator;
+/// # use miniscript::{Descriptor, DescriptorPublicKey};
+/// # use bitcoin::{secp256k1::Secp256k1};
+/// # use std::str::FromStr;
+/// # let secp = bitcoin::secp256k1::Secp256k1::signing_only();
+/// # let (descriptor, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "wpkh([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/0)").unwrap();
+/// # let external_spk_0 = descriptor.at_derivation_index(0).script_pubkey();
+/// # let external_spk_3 = descriptor.at_derivation_index(3).script_pubkey();
+/// # let external_spk_4 = descriptor.at_derivation_index(4).script_pubkey();
+///
+/// // Creates a new script pubkey iterator starting at 0 from a descriptor.
+/// let mut spk_iter = SpkIterator::new(&descriptor);
+/// assert_eq!(spk_iter.next(), Some((0, external_spk_0)));
+/// assert_eq!(spk_iter.next(), None);
+/// ```
+#[derive(Clone)]
+pub struct SpkIterator<D> {
+    next_index: u32,
+    end: u32,
+    descriptor: D,
+    secp: Secp256k1<bitcoin::secp256k1::VerifyOnly>,
+}
+
+impl<D> SpkIterator<D>
+where
+    D: Borrow<Descriptor<DescriptorPublicKey>>,
+{
+    /// Creates a new script pubkey iterator starting at 0 from a descriptor.
+    pub fn new(descriptor: D) -> Self {
+        let end = if descriptor.borrow().has_wildcard() {
+            BIP32_MAX_INDEX
+        } else {
+            0
+        };
+
+        SpkIterator::new_with_range(descriptor, 0..=end)
+    }
+
+    // Creates a new script pubkey iterator from a descriptor with a given range.
+    pub(crate) fn new_with_range<R>(descriptor: D, range: R) -> Self
+    where
+        R: RangeBounds<u32>,
+    {
+        let mut end = match range.end_bound() {
+            Bound::Included(end) => *end + 1,
+            Bound::Excluded(end) => *end,
+            Bound::Unbounded => u32::MAX,
+        };
+        // Because `end` is exclusive, we want the maximum value to be BIP32_MAX_INDEX + 1.
+        end = end.min(BIP32_MAX_INDEX + 1);
+
+        Self {
+            next_index: match range.start_bound() {
+                Bound::Included(start) => *start,
+                Bound::Excluded(start) => *start + 1,
+                Bound::Unbounded => u32::MIN,
+            },
+            end,
+            descriptor,
+            secp: Secp256k1::verification_only(),
+        }
+    }
+}
+
+impl<D> Iterator for SpkIterator<D>
+where
+    D: Borrow<Descriptor<DescriptorPublicKey>>,
+{
+    type Item = (u32, Script);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // For non-wildcard descriptors, we expect the first element to be Some((0, spk)), then None after.
+        // For wildcard descriptors, we expect it to keep iterating until exhausted.
+        if self.next_index >= self.end {
+            return None;
+        }
+
+        let script = self
+            .descriptor
+            .borrow()
+            .at_derivation_index(self.next_index)
+            .derived_descriptor(&self.secp)
+            .expect("the descriptor cannot need hardened derivation")
+            .script_pubkey();
+        let output = (self.next_index, script);
+
+        self.next_index += 1;
+
+        Some(output)
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.next_index = self
+            .next_index
+            .saturating_add(u32::try_from(n).unwrap_or(u32::MAX));
+        self.next()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::{
+        bitcoin::secp256k1::Secp256k1,
+        keychain::KeychainTxOutIndex,
+        miniscript::{Descriptor, DescriptorPublicKey},
+        spk_iter::{SpkIterator, BIP32_MAX_INDEX},
+    };
+
+    #[derive(Clone, Debug, PartialEq, Eq, Ord, PartialOrd)]
+    enum TestKeychain {
+        External,
+        Internal,
+    }
+
+    fn init_txout_index() -> (
+        KeychainTxOutIndex<TestKeychain>,
+        Descriptor<DescriptorPublicKey>,
+        Descriptor<DescriptorPublicKey>,
+    ) {
+        let mut txout_index = KeychainTxOutIndex::<TestKeychain>::default();
+
+        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(TestKeychain::External, external_descriptor.clone());
+        txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone());
+
+        (txout_index, external_descriptor, internal_descriptor)
+    }
+
+    #[test]
+    #[allow(clippy::iter_nth_zero)]
+    fn test_spkiterator_wildcard() {
+        let (_, external_desc, _) = init_txout_index();
+        let external_spk_0 = external_desc.at_derivation_index(0).script_pubkey();
+        let external_spk_16 = external_desc.at_derivation_index(16).script_pubkey();
+        let external_spk_20 = external_desc.at_derivation_index(20).script_pubkey();
+        let external_spk_21 = external_desc.at_derivation_index(21).script_pubkey();
+        let external_spk_max = external_desc
+            .at_derivation_index(BIP32_MAX_INDEX)
+            .script_pubkey();
+
+        let mut external_spk = SpkIterator::new(&external_desc);
+        let max_index = BIP32_MAX_INDEX - 22;
+
+        assert_eq!(external_spk.next().unwrap(), (0, external_spk_0));
+        assert_eq!(external_spk.nth(15).unwrap(), (16, external_spk_16));
+        assert_eq!(external_spk.nth(3).unwrap(), (20, external_spk_20.clone()));
+        assert_eq!(external_spk.next().unwrap(), (21, external_spk_21));
+        assert_eq!(
+            external_spk.nth(max_index as usize).unwrap(),
+            (BIP32_MAX_INDEX, external_spk_max)
+        );
+        assert_eq!(external_spk.nth(0), None);
+
+        let mut external_spk = SpkIterator::new_with_range(&external_desc, 0..21);
+        assert_eq!(external_spk.nth(20).unwrap(), (20, external_spk_20));
+        assert_eq!(external_spk.next(), None);
+
+        let mut external_spk = SpkIterator::new_with_range(&external_desc, 0..21);
+        assert_eq!(external_spk.nth(21), None);
+    }
+
+    #[test]
+    #[allow(clippy::iter_nth_zero)]
+    fn test_spkiterator_non_wildcard() {
+        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();
+        let external_spk_0 = no_wildcard_descriptor
+            .at_derivation_index(0)
+            .script_pubkey();
+
+        let mut external_spk = SpkIterator::new(&no_wildcard_descriptor);
+
+        assert_eq!(external_spk.next().unwrap(), (0, external_spk_0.clone()));
+        assert_eq!(external_spk.next(), None);
+
+        let mut external_spk = SpkIterator::new(&no_wildcard_descriptor);
+
+        assert_eq!(external_spk.nth(0).unwrap(), (0, external_spk_0));
+        assert_eq!(external_spk.nth(0), None);
+    }
+
+    // The following dummy traits were created to test if SpkIterator is working properly.
+    trait TestSendStatic: Send + 'static {
+        fn test(&self) -> u32 {
+            20
+        }
+    }
+
+    impl TestSendStatic for SpkIterator<Descriptor<DescriptorPublicKey>> {
+        fn test(&self) -> u32 {
+            20
+        }
+    }
+}