Skip to content

Commit

Permalink
Updated custom spk iterator for PR#927
Browse files Browse the repository at this point in the history
  • Loading branch information
@LagginTimes
LagginTimes committed Apr 15, 2023
1 parent 091b068 commit 979ecdd
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 29 deletions.
59 changes: 31 additions & 28 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, ops::Range};
use core::{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..BIP32_MAX_INDEX + 1),
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..BIP32_MAX_INDEX + 1)
SpkIterator::new_with_range(Cow::Owned(descriptor), 0..)
}

/// Convenience method to get [`revealed_spks_of_keychain`] of all keychains.
Expand Down Expand Up @@ -400,8 +400,8 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
}

// we range over indexes that are not stored
let range = next_reveal_index + lookahead..target_index + lookahead + 1;
for (new_index, new_spk) in range_descriptor_spks(Cow::Borrowed(descriptor), range) {
let range = next_reveal_index + lookahead..=target_index + lookahead;
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,18 +559,6 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
}
}

fn range_descriptor_spks(
descriptor: Cow<'_, Descriptor<DescriptorPublicKey>>,
range: Range<u32>,
) -> SpkIterator {
SpkIterator {
next_index: range.start,
end: range.end,
descriptor,
secp: Secp256k1::verification_only(),
}
}

#[derive(Clone)]
pub struct SpkIterator<'a> {
next_index: u32,
Expand All @@ -582,18 +570,33 @@ pub struct SpkIterator<'a> {
impl<'a> SpkIterator<'a> {
// Creates a new script pubkey iterator starting at 0 from a descriptor
pub fn new(descriptor: Cow<'a, Descriptor<DescriptorPublicKey>>) -> Self {
let secp = Secp256k1::verification_only();
let end = if descriptor.has_wildcard() {
BIP32_MAX_INDEX
} else {
0
};

SpkIterator::new_with_range(descriptor, 0..=end)
}

// Creates a new script pubkey iterator starting at 0 from a descriptor with a given range
pub fn new_with_range<R>(descriptor: Cow<'a, Descriptor<DescriptorPublicKey>>, range: R) -> Self
where
R: RangeBounds<u32>,
{
Self {
next_index: 0,
end,
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,
},
descriptor,
secp,
secp: Secp256k1::verification_only(),
}
}
}
Expand All @@ -607,7 +610,7 @@ impl<'a> Iterator for SpkIterator<'a> {
// for wildcard descriptors:
// * we expect it to keep iterating until exhausted.
let has_wildcard = self.descriptor.has_wildcard();
if self.next_index >= self.end || (has_wildcard && self.next_index >= BIP32_MAX_INDEX) {
if self.next_index >= self.end || (has_wildcard && self.next_index > BIP32_MAX_INDEX) {
return None;
}

Expand All @@ -625,9 +628,9 @@ impl<'a> Iterator for SpkIterator<'a> {
}

fn nth(&mut self, n: usize) -> Option<Self::Item> {
//let n = self.next_index + n as u32;
//self.descriptor.at_derivation_index(n);
self.next_index = self.next_index.saturating_add(n as u32);
self.next_index = self
.next_index
.saturating_add(u32::try_from(n).unwrap_or(u32::MAX));
self.next()
}
}
32 changes: 31 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,7 +4,7 @@
mod common;
use bdk_chain::{
collections::BTreeMap,
keychain::{DerivationAdditions, KeychainTxOutIndex, SpkIterator},
keychain::{DerivationAdditions, KeychainTxOutIndex, SpkIterator, BIP32_MAX_INDEX},
};

use bitcoin::{secp256k1::Secp256k1, OutPoint, Script, Transaction, TxOut};
Expand Down Expand Up @@ -370,17 +370,47 @@ fn test_non_wildcard_derivations() {
}

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

#[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 979ecdd

Please sign in to comment.