Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Custom spk iterator #927

Merged
merged 1 commit into from
May 2, 2023
Merged

Custom spk iterator #927

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 21 additions & 54 deletions crates/chain/src/keychain/txout_index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down Expand Up @@ -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(
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 @@ -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()
Expand All @@ -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.
Expand Down Expand Up @@ -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();
Expand All @@ -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());
}
}

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

Expand All @@ -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);
Expand All @@ -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(
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(
LagginTimes marked this conversation as resolved.
Show resolved Hide resolved
descriptor.clone(),
next_reveal_index..next_reveal_index,
),
DerivationAdditions::default(),
Expand Down Expand Up @@ -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)
}
4 changes: 4 additions & 0 deletions crates/chain/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
215 changes: 215 additions & 0 deletions crates/chain/src/spk_iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
}
}