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 and evanlinjin committed Apr 20, 2023
1 parent 82f9cad commit 10fb618
Show file tree
Hide file tree
Showing 3 changed files with 230 additions and 50 deletions.
70 changes: 20 additions & 50 deletions crates/chain/src/keychain/txout_index.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
use crate::{
collections::*,
miniscript::{Descriptor, DescriptorPublicKey},
spk_iter::SpkIterator,
ForEachTxOut, 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 super::DerivationAdditions;
Expand Down Expand Up @@ -214,10 +215,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(
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);
Expand All @@ -237,13 +237,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()
Expand All @@ -255,13 +255,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.
Expand Down Expand Up @@ -341,7 +341,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();
Expand All @@ -351,7 +351,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());
}
}

Expand All @@ -376,7 +376,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();

Expand All @@ -401,7 +404,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);
Expand All @@ -418,16 +421,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(
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(),
Expand Down Expand Up @@ -558,33 +558,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)
}
1 change: 1 addition & 0 deletions crates/chain/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ pub mod sparse_chain;
mod tx_data_traits;
pub mod tx_graph;
pub use tx_data_traits::*;
pub mod spk_iter;

#[doc(hidden)]
pub mod example_utils;
Expand Down
209 changes: 209 additions & 0 deletions crates/chain/src/spk_iter.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,209 @@
use crate::{
bitcoin::{secp256k1::Secp256k1, Script},
keychain::BIP32_MAX_INDEX,
miniscript::{Descriptor, DescriptorPublicKey},
};
use core::{borrow::Borrow, ops::Bound, ops::RangeBounds};

/// 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::spk_iter::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().unwrap(), (0, external_spk_0));
/// assert_eq!(spk_iter.next(), None);
///
/// // 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 mut spk_range_iter = SpkIterator::new_with_range(&descriptor, 0..BIP32_MAX_INDEX);
/// assert_eq!(spk_range_iter.nth(3).unwrap(), (3, external_spk_3));
/// assert_eq!(spk_range_iter.next().unwrap(), (4, external_spk_4));
///
/// ```
#[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 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.
if end > BIP32_MAX_INDEX + 1 {
end = 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, BIP32_MAX_INDEX},
miniscript::{Descriptor, DescriptorPublicKey},
spk_iter::SpkIterator,
};

#[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);
}
}

0 comments on commit 10fb618

Please sign in to comment.