Skip to content

Commit

Permalink
Implement SpkIterator
Browse files Browse the repository at this point in the history 
SpkIterator was created with its own nth() and next() implementations
and its own new() and new_with_range() constructors.

Co-authored-by: 志宇 <[email protected]>
  • Loading branch information
@LagginTimes @evanlinjin
LagginTimes and evanlinjin committed Apr 19, 2023
1 parent 82f9cad commit 1217917
Show file tree
Hide file tree
Showing 2 changed files with 165 additions and 33 deletions.
142 changes: 110 additions & 32 deletions crates/chain/src/keychain/txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -214,7 +214,7 @@ 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 SpkIterator::new_with_range(
Cow::Borrowed(descriptor),
next_store_index..next_reveal_index + lookahead,
) {
Expand Down Expand Up @@ -243,7 +243,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
.map(|(keychain, descriptor)| {
(
keychain.clone(),
range_descriptor_spks(Cow::Owned(descriptor.clone()), 0..),
SpkIterator::new_with_range(Cow::Owned(descriptor.clone()), 0..),
)
})
.collect()
Expand All @@ -261,7 +261,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
.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.
Expand Down Expand Up @@ -401,7 +401,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(Cow::Borrowed(descriptor), range) {
let _inserted = self
.inner
.insert_spk((keychain.clone(), new_index), new_spk);
Expand All @@ -418,15 +418,15 @@ 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(
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()),
next_reveal_index..next_reveal_index,
),
Expand Down Expand Up @@ -559,32 +559,110 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
}
}

fn range_descriptor_spks<'a, R>(
descriptor: Cow<'a, Descriptor<DescriptorPublicKey>>,
range: R,
) -> impl Iterator<Item = (u32, Script)> + 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::<DescriptorPublicKey>::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<D> {
next_index: u32,
end: u32,
descriptor: D,
secp: Secp256k1<bitcoin::secp256k1::VerifyOnly>,
}

impl<D> SpkIterator<D>
where
R: Iterator<Item = u32> + Clone + Send + 'a,
D: Borrow<Descriptor<DescriptorPublicKey>>,
{
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<R>(descriptor: D, range: R) -> Self
where
R: RangeBounds<u32>,
{
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<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 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::Item> {
self.next_index = self
.next_index
.saturating_add(u32::try_from(n).unwrap_or(u32::MAX));
self.next()
}
}
56 changes: 55 additions & 1 deletion crates/chain/tests/test_keychain_txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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::<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(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);
}

0 comments on commit 1217917

Please sign in to comment.