|
| 1 | +use crate::{ |
| 2 | + bitcoin::{secp256k1::Secp256k1, Script}, |
| 3 | + miniscript::{Descriptor, DescriptorPublicKey}, |
| 4 | +}; |
| 5 | +use core::{borrow::Borrow, ops::Bound, ops::RangeBounds}; |
| 6 | + |
| 7 | +/// Maximum [BIP32](https://bips.xyz/32) derivation index. |
| 8 | +pub const BIP32_MAX_INDEX: u32 = (1 << 31) - 1; |
| 9 | + |
| 10 | +/// An iterator for derived script pubkeys. |
| 11 | +/// |
| 12 | +/// [`SpkIterator`] is an implementation of the [`Iterator`] trait which possesses its own `next()` |
| 13 | +/// and `nth()` functions, both of which circumvent the unnecessary intermediate derivations required |
| 14 | +/// when using their default implementations. |
| 15 | +/// |
| 16 | +/// ## Examples |
| 17 | +/// |
| 18 | +/// ``` |
| 19 | +/// use bdk_chain::SpkIterator; |
| 20 | +/// # use miniscript::{Descriptor, DescriptorPublicKey}; |
| 21 | +/// # use bitcoin::{secp256k1::Secp256k1}; |
| 22 | +/// # use std::str::FromStr; |
| 23 | +/// # let secp = bitcoin::secp256k1::Secp256k1::signing_only(); |
| 24 | +/// # let (descriptor, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "wpkh([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/0)").unwrap(); |
| 25 | +/// # let external_spk_0 = descriptor.at_derivation_index(0).script_pubkey(); |
| 26 | +/// # let external_spk_3 = descriptor.at_derivation_index(3).script_pubkey(); |
| 27 | +/// # let external_spk_4 = descriptor.at_derivation_index(4).script_pubkey(); |
| 28 | +/// |
| 29 | +/// // Creates a new script pubkey iterator starting at 0 from a descriptor. |
| 30 | +/// let mut spk_iter = SpkIterator::new(&descriptor); |
| 31 | +/// assert_eq!(spk_iter.next(), Some((0, external_spk_0))); |
| 32 | +/// assert_eq!(spk_iter.next(), None); |
| 33 | +/// ``` |
| 34 | +#[derive(Clone)] |
| 35 | +pub struct SpkIterator<D> { |
| 36 | + next_index: u32, |
| 37 | + end: u32, |
| 38 | + descriptor: D, |
| 39 | + secp: Secp256k1<bitcoin::secp256k1::VerifyOnly>, |
| 40 | +} |
| 41 | + |
| 42 | +impl<D> SpkIterator<D> |
| 43 | +where |
| 44 | + D: Borrow<Descriptor<DescriptorPublicKey>>, |
| 45 | +{ |
| 46 | + /// Creates a new script pubkey iterator starting at 0 from a descriptor. |
| 47 | + pub fn new(descriptor: D) -> Self { |
| 48 | + let end = if descriptor.borrow().has_wildcard() { |
| 49 | + BIP32_MAX_INDEX |
| 50 | + } else { |
| 51 | + 0 |
| 52 | + }; |
| 53 | + |
| 54 | + SpkIterator::new_with_range(descriptor, 0..=end) |
| 55 | + } |
| 56 | + |
| 57 | + // Creates a new script pubkey iterator from a descriptor with a given range. |
| 58 | + pub(crate) fn new_with_range<R>(descriptor: D, range: R) -> Self |
| 59 | + where |
| 60 | + R: RangeBounds<u32>, |
| 61 | + { |
| 62 | + let mut end = match range.end_bound() { |
| 63 | + Bound::Included(end) => *end + 1, |
| 64 | + Bound::Excluded(end) => *end, |
| 65 | + Bound::Unbounded => u32::MAX, |
| 66 | + }; |
| 67 | + // Because `end` is exclusive, we want the maximum value to be BIP32_MAX_INDEX + 1. |
| 68 | + end = end.min(BIP32_MAX_INDEX + 1); |
| 69 | + |
| 70 | + Self { |
| 71 | + next_index: match range.start_bound() { |
| 72 | + Bound::Included(start) => *start, |
| 73 | + Bound::Excluded(start) => *start + 1, |
| 74 | + Bound::Unbounded => u32::MIN, |
| 75 | + }, |
| 76 | + end, |
| 77 | + descriptor, |
| 78 | + secp: Secp256k1::verification_only(), |
| 79 | + } |
| 80 | + } |
| 81 | +} |
| 82 | + |
| 83 | +impl<D> Iterator for SpkIterator<D> |
| 84 | +where |
| 85 | + D: Borrow<Descriptor<DescriptorPublicKey>>, |
| 86 | +{ |
| 87 | + type Item = (u32, Script); |
| 88 | + |
| 89 | + fn next(&mut self) -> Option<Self::Item> { |
| 90 | + // For non-wildcard descriptors, we expect the first element to be Some((0, spk)), then None after. |
| 91 | + // For wildcard descriptors, we expect it to keep iterating until exhausted. |
| 92 | + if self.next_index >= self.end { |
| 93 | + return None; |
| 94 | + } |
| 95 | + |
| 96 | + let script = self |
| 97 | + .descriptor |
| 98 | + .borrow() |
| 99 | + .at_derivation_index(self.next_index) |
| 100 | + .derived_descriptor(&self.secp) |
| 101 | + .expect("the descriptor cannot need hardened derivation") |
| 102 | + .script_pubkey(); |
| 103 | + let output = (self.next_index, script); |
| 104 | + |
| 105 | + self.next_index += 1; |
| 106 | + |
| 107 | + Some(output) |
| 108 | + } |
| 109 | + |
| 110 | + fn nth(&mut self, n: usize) -> Option<Self::Item> { |
| 111 | + self.next_index = self |
| 112 | + .next_index |
| 113 | + .saturating_add(u32::try_from(n).unwrap_or(u32::MAX)); |
| 114 | + self.next() |
| 115 | + } |
| 116 | +} |
| 117 | + |
| 118 | +#[cfg(test)] |
| 119 | +mod test { |
| 120 | + use crate::{ |
| 121 | + bitcoin::secp256k1::Secp256k1, |
| 122 | + keychain::KeychainTxOutIndex, |
| 123 | + miniscript::{Descriptor, DescriptorPublicKey}, |
| 124 | + spk_iter::{SpkIterator, BIP32_MAX_INDEX}, |
| 125 | + }; |
| 126 | + |
| 127 | + #[derive(Clone, Debug, PartialEq, Eq, Ord, PartialOrd)] |
| 128 | + enum TestKeychain { |
| 129 | + External, |
| 130 | + Internal, |
| 131 | + } |
| 132 | + |
| 133 | + fn init_txout_index() -> ( |
| 134 | + KeychainTxOutIndex<TestKeychain>, |
| 135 | + Descriptor<DescriptorPublicKey>, |
| 136 | + Descriptor<DescriptorPublicKey>, |
| 137 | + ) { |
| 138 | + let mut txout_index = KeychainTxOutIndex::<TestKeychain>::default(); |
| 139 | + |
| 140 | + let secp = Secp256k1::signing_only(); |
| 141 | + let (external_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap(); |
| 142 | + let (internal_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/*)").unwrap(); |
| 143 | + |
| 144 | + txout_index.add_keychain(TestKeychain::External, external_descriptor.clone()); |
| 145 | + txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone()); |
| 146 | + |
| 147 | + (txout_index, external_descriptor, internal_descriptor) |
| 148 | + } |
| 149 | + |
| 150 | + #[test] |
| 151 | + #[allow(clippy::iter_nth_zero)] |
| 152 | + fn test_spkiterator_wildcard() { |
| 153 | + let (_, external_desc, _) = init_txout_index(); |
| 154 | + let external_spk_0 = external_desc.at_derivation_index(0).script_pubkey(); |
| 155 | + let external_spk_16 = external_desc.at_derivation_index(16).script_pubkey(); |
| 156 | + let external_spk_20 = external_desc.at_derivation_index(20).script_pubkey(); |
| 157 | + let external_spk_21 = external_desc.at_derivation_index(21).script_pubkey(); |
| 158 | + let external_spk_max = external_desc |
| 159 | + .at_derivation_index(BIP32_MAX_INDEX) |
| 160 | + .script_pubkey(); |
| 161 | + |
| 162 | + let mut external_spk = SpkIterator::new(&external_desc); |
| 163 | + let max_index = BIP32_MAX_INDEX - 22; |
| 164 | + |
| 165 | + assert_eq!(external_spk.next().unwrap(), (0, external_spk_0)); |
| 166 | + assert_eq!(external_spk.nth(15).unwrap(), (16, external_spk_16)); |
| 167 | + assert_eq!(external_spk.nth(3).unwrap(), (20, external_spk_20.clone())); |
| 168 | + assert_eq!(external_spk.next().unwrap(), (21, external_spk_21)); |
| 169 | + assert_eq!( |
| 170 | + external_spk.nth(max_index as usize).unwrap(), |
| 171 | + (BIP32_MAX_INDEX, external_spk_max) |
| 172 | + ); |
| 173 | + assert_eq!(external_spk.nth(0), None); |
| 174 | + |
| 175 | + let mut external_spk = SpkIterator::new_with_range(&external_desc, 0..21); |
| 176 | + assert_eq!(external_spk.nth(20).unwrap(), (20, external_spk_20)); |
| 177 | + assert_eq!(external_spk.next(), None); |
| 178 | + |
| 179 | + let mut external_spk = SpkIterator::new_with_range(&external_desc, 0..21); |
| 180 | + assert_eq!(external_spk.nth(21), None); |
| 181 | + } |
| 182 | + |
| 183 | + #[test] |
| 184 | + #[allow(clippy::iter_nth_zero)] |
| 185 | + fn test_spkiterator_non_wildcard() { |
| 186 | + let secp = bitcoin::secp256k1::Secp256k1::signing_only(); |
| 187 | + let (no_wildcard_descriptor, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "wpkh([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/0)").unwrap(); |
| 188 | + let external_spk_0 = no_wildcard_descriptor |
| 189 | + .at_derivation_index(0) |
| 190 | + .script_pubkey(); |
| 191 | + |
| 192 | + let mut external_spk = SpkIterator::new(&no_wildcard_descriptor); |
| 193 | + |
| 194 | + assert_eq!(external_spk.next().unwrap(), (0, external_spk_0.clone())); |
| 195 | + assert_eq!(external_spk.next(), None); |
| 196 | + |
| 197 | + let mut external_spk = SpkIterator::new(&no_wildcard_descriptor); |
| 198 | + |
| 199 | + assert_eq!(external_spk.nth(0).unwrap(), (0, external_spk_0)); |
| 200 | + assert_eq!(external_spk.nth(0), None); |
| 201 | + } |
| 202 | + |
| 203 | + // The following dummy traits were created to test if SpkIterator is working properly. |
| 204 | + trait TestSendStatic: Send + 'static { |
| 205 | + fn test(&self) -> u32 { |
| 206 | + 20 |
| 207 | + } |
| 208 | + } |
| 209 | + |
| 210 | + impl TestSendStatic for SpkIterator<Descriptor<DescriptorPublicKey>> { |
| 211 | + fn test(&self) -> u32 { |
| 212 | + 20 |
| 213 | + } |
| 214 | + } |
| 215 | +} |
0 commit comments