frost/bdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement SpkIterator

Browse Source 
SpkIterator was created with its own nth() and next() implementations
and its own new() and new_with_range() constructors.

Co-authored-by: 志宇 <hello@evanlinjin.me>
This commit is contained in:
Wei 2023-03-22 17:00:08 +08:00 committed by 志宇
parent 139e3d3802
commit 10fe32e6f1
No known key found for this signature in database
GPG Key ID: F6345C9837C2BDE8
3 changed files with 240 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
crates/chain/src/keychain/txout_index.rs
View File

@ -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<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);
@ -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()
@ -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.
@ -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();
@ -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());
}
}
@ -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();
@ -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);
@ -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(
descriptor.clone(),
next_reveal_index..next_reveal_index,
),
DerivationAdditions::default(),
@ -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
crates/chain/src/lib.rs
View File

@ -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]

215
crates/chain/src/spk_iter.rs Normal file
View File

@ -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
}
}
}