From 10fe32e6f10091c0c43aba3f66e0f700a06f8c89 Mon Sep 17 00:00:00 2001 From: Wei <122563728+LagginTimes@users.noreply.github.com> Date: Wed, 22 Mar 2023 17:00:08 +0800 Subject: [PATCH] Implement SpkIterator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SpkIterator was created with its own nth() and next() implementations and its own new() and new_with_range() constructors. Co-authored-by: 志宇 --- crates/chain/src/keychain/txout_index.rs | 75 +++----- crates/chain/src/lib.rs | 4 + crates/chain/src/spk_iter.rs | 215 +++++++++++++++++++++++ 3 files changed, 240 insertions(+), 54 deletions(-) create mode 100644 crates/chain/src/spk_iter.rs diff --git a/crates/chain/src/keychain/txout_index.rs b/crates/chain/src/keychain/txout_index.rs index fbe67d1f..f8ef46be 100644 --- a/crates/chain/src/keychain/txout_index.rs +++ b/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 bitcoin::{secp256k1::Secp256k1, OutPoint, Script, TxOut}; +use alloc::vec::Vec; +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 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( - Cow::Borrowed(descriptor), - next_store_index..next_reveal_index + lookahead, - ) { + for (new_index, new_spk) in + SpkIterator::new_with_range(descriptor, next_store_index..next_reveal_index + lookahead) + { let _inserted = self .inner .insert_spk((keychain.clone(), new_index), new_spk); @@ -266,13 +263,13 @@ impl KeychainTxOutIndex { /// derivable script pubkeys. pub fn spks_of_all_keychains( &self, - ) -> BTreeMap + Clone> { + ) -> BTreeMap>> { 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 KeychainTxOutIndex { /// # Panics /// /// This will panic if the `keychain` does not exist. - pub fn spks_of_keychain(&self, keychain: &K) -> impl Iterator + Clone { + pub fn spks_of_keychain(&self, keychain: &K) -> SpkIterator> { 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 KeychainTxOutIndex { &mut self, keychains: &BTreeMap, ) -> ( - BTreeMap>, + BTreeMap>>, DerivationAdditions, ) { let mut additions = DerivationAdditions::default(); @@ -380,7 +377,7 @@ impl KeychainTxOutIndex { 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 KeychainTxOutIndex { &mut self, keychain: &K, target_index: u32, - ) -> (impl Iterator, DerivationAdditions) { + ) -> ( + SpkIterator>, + DerivationAdditions, + ) { let descriptor = self.keychains.get(keychain).expect("keychain must exist"); let has_wildcard = descriptor.has_wildcard(); @@ -430,7 +430,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(descriptor, range) { let _inserted = self .inner .insert_spk((keychain.clone(), new_index), new_spk); @@ -447,16 +447,13 @@ impl KeychainTxOutIndex { let _old_index = self.last_revealed.insert(keychain.clone(), index); debug_assert!(_old_index < Some(index)); ( - range_descriptor_spks( - Cow::Owned(descriptor.clone()), - next_reveal_index..index + 1, - ), + SpkIterator::new_with_range(descriptor.clone(), next_reveal_index..index + 1), DerivationAdditions(core::iter::once((keychain.clone(), index)).collect()), ) } None => ( - range_descriptor_spks( - Cow::Owned(descriptor.clone()), + SpkIterator::new_with_range( + descriptor.clone(), next_reveal_index..next_reveal_index, ), DerivationAdditions::default(), @@ -587,33 +584,3 @@ impl KeychainTxOutIndex { let _ = self.reveal_to_target_multi(&additions.0); } } - -fn range_descriptor_spks<'a, R>( - descriptor: Cow<'a, Descriptor>, - range: R, -) -> impl Iterator + Clone + Send + 'a -where - R: Iterator + 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) -} diff --git a/crates/chain/src/lib.rs b/crates/chain/src/lib.rs index 26527623..cf3cda3b 100644 --- a/crates/chain/src/lib.rs +++ b/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] diff --git a/crates/chain/src/spk_iter.rs b/crates/chain/src/spk_iter.rs new file mode 100644 index 00000000..97c81441 --- /dev/null +++ b/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::::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 { + next_index: u32, + end: u32, + descriptor: D, + secp: Secp256k1, +} + +impl SpkIterator +where + D: Borrow>, +{ + /// 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(descriptor: D, range: R) -> Self + where + R: RangeBounds, + { + 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 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 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.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, + Descriptor, + Descriptor, + ) { + let mut txout_index = KeychainTxOutIndex::::default(); + + let secp = Secp256k1::signing_only(); + let (external_descriptor,_) = Descriptor::::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap(); + let (internal_descriptor,_) = Descriptor::::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::::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> { + fn test(&self) -> u32 { + 20 + } + } +}