diff --git a/crates/chain/src/keychain/txout_index.rs b/crates/chain/src/keychain/txout_index.rs index feb71edb45..dce1286708 100644 --- a/crates/chain/src/keychain/txout_index.rs +++ b/crates/chain/src/keychain/txout_index.rs @@ -5,7 +5,7 @@ use crate::{ }; use alloc::{borrow::Cow, vec::Vec}; use bitcoin::{secp256k1::Secp256k1, OutPoint, Script, TxOut}; -use core::{fmt::Debug, ops::Deref}; +use core::{borrow::Borrow, fmt::Debug, ops::Bound, ops::Deref, ops::RangeBounds}; use super::DerivationAdditions; @@ -214,7 +214,7 @@ impl KeychainTxOutIndex { 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 SpkIterator::new_with_range( Cow::Borrowed(descriptor), next_store_index..next_reveal_index + lookahead, ) { @@ -243,7 +243,7 @@ impl KeychainTxOutIndex { .map(|(keychain, descriptor)| { ( keychain.clone(), - range_descriptor_spks(Cow::Owned(descriptor.clone()), 0..), + SpkIterator::new_with_range(Cow::Owned(descriptor.clone()), 0..), ) }) .collect() @@ -261,7 +261,7 @@ impl KeychainTxOutIndex { .get(keychain) .expect("keychain must exist") .clone(); - range_descriptor_spks(Cow::Owned(descriptor), 0..) + SpkIterator::new_with_range(Cow::Owned(descriptor), 0..) } /// Convenience method to get [`revealed_spks_of_keychain`] of all keychains. @@ -401,7 +401,7 @@ impl KeychainTxOutIndex { // 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(Cow::Borrowed(descriptor), range) { let _inserted = self .inner .insert_spk((keychain.clone(), new_index), new_spk); @@ -418,7 +418,7 @@ impl KeychainTxOutIndex { let _old_index = self.last_revealed.insert(keychain.clone(), index); debug_assert!(_old_index < Some(index)); ( - range_descriptor_spks( + SpkIterator::new_with_range( Cow::Owned(descriptor.clone()), next_reveal_index..index + 1, ), @@ -426,7 +426,7 @@ impl KeychainTxOutIndex { ) } None => ( - range_descriptor_spks( + SpkIterator::new_with_range( Cow::Owned(descriptor.clone()), next_reveal_index..next_reveal_index, ), @@ -559,32 +559,110 @@ impl KeychainTxOutIndex { } } -fn range_descriptor_spks<'a, R>( - descriptor: Cow<'a, Descriptor>, - range: R, -) -> impl Iterator + Clone + Send + 'a +/// 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::keychain::SpkIterator; +/// use std::borrow::Cow; +/// # use miniscript::{Descriptor, DescriptorPublicKey}; +/// # use bitcoin::{secp256k1::Secp256k1}; +/// # use std::str::FromStr; +/// # let secp = bitcoin::secp256k1::Secp256k1::signing_only(); +/// # let (descriptor, _) = Descriptor::::parse_descriptor(&secp, "wpkh([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/0)").unwrap(); +/// +/// // Creates a new script pubkey iterator starting at 0 from a descriptor. +/// let _spk_iter = SpkIterator::new(Cow::Borrowed(&descriptor)); +/// +/// // Creates a new script pubkey iterator from a descriptor with a given range, in this +/// // case between 0 and a constant representing the maximum BIP32 derivation index. +/// const BIP32_MAX_INDEX: u32 = (1 << 31) - 1; +/// let _spk_range_iter = SpkIterator::new_with_range(Cow::Borrowed(&descriptor), 0..BIP32_MAX_INDEX); +/// ``` +#[derive(Clone)] +pub struct SpkIterator { + next_index: u32, + end: u32, + descriptor: D, + secp: Secp256k1, +} + +impl SpkIterator where - R: Iterator + Clone + Send + 'a, + D: Borrow>, { - 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(), - )) + /// 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 fn new_with_range(descriptor: D, range: R) -> Self + where + R: RangeBounds, + { + Self { + next_index: match range.start_bound() { + Bound::Included(start) => *start, + Bound::Excluded(start) => *start + 1, + Bound::Unbounded => u32::MIN, + }, + end: match range.end_bound() { + Bound::Included(end) => *end + 1, + Bound::Excluded(end) => *end, + Bound::Unbounded => u32::MAX, }, - ) - .take_while(Result::is_ok) - .map(Result::unwrap) + descriptor, + secp: Secp256k1::verification_only(), + } + } +} + +impl Iterator for SpkIterator +where + D: Borrow>, +{ + type Item = (u32, Script); + + fn next(&mut self) -> Option { + // for non-wildcard descriptors: + // * we expect the first element to be Some((0, spk)), then will return None after. + // for wildcard descriptors: + // * we expect it to keep iterating until exhausted. + let has_wildcard = self.descriptor.borrow().has_wildcard(); + if self.next_index >= self.end || (has_wildcard && self.next_index > BIP32_MAX_INDEX) { + 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.next_index = self + .next_index + .saturating_add(u32::try_from(n).unwrap_or(u32::MAX)); + self.next() + } } diff --git a/crates/chain/tests/test_keychain_txout_index.rs b/crates/chain/tests/test_keychain_txout_index.rs index 07c7f48d4d..8262ebea13 100644 --- a/crates/chain/tests/test_keychain_txout_index.rs +++ b/crates/chain/tests/test_keychain_txout_index.rs @@ -4,11 +4,12 @@ mod common; use bdk_chain::{ collections::BTreeMap, - keychain::{DerivationAdditions, KeychainTxOutIndex}, + keychain::{DerivationAdditions, KeychainTxOutIndex, SpkIterator, BIP32_MAX_INDEX}, }; use bitcoin::{secp256k1::Secp256k1, OutPoint, Script, Transaction, TxOut}; use miniscript::{Descriptor, DescriptorPublicKey}; +use std::borrow::Cow; #[derive(Clone, Debug, PartialEq, Eq, Ord, PartialOrd)] enum TestKeychain { @@ -367,3 +368,56 @@ fn test_non_wildcard_derivations() { assert_eq!(revealed_spks.count(), 0); assert!(revealed_additions.is_empty()); } + +#[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(Cow::Borrowed(&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(Cow::Borrowed(&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(Cow::Owned(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::::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(Cow::Borrowed(&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(Cow::Owned(no_wildcard_descriptor)); + + assert_eq!(external_spk.nth(0).unwrap(), (0, external_spk_0)); + assert_eq!(external_spk.nth(0), None); +}