From 523654e78e6ef093211fadf02cbd96823d169801 Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Fri, 23 May 2025 23:03:51 +0000 Subject: [PATCH 1/2] descriptor: add Descriptor::iter_pk This commit was -not- AI-generated, though probably it should've been. --- src/descriptor/iter.rs | 127 +++++++++++++++++++++++++++++++++++++++++ src/descriptor/mod.rs | 25 ++++++++ 2 files changed, 152 insertions(+) create mode 100644 src/descriptor/iter.rs diff --git a/src/descriptor/iter.rs b/src/descriptor/iter.rs new file mode 100644 index 000000000..3e5552b4c --- /dev/null +++ b/src/descriptor/iter.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: CC0-1.0 + +//! Iterators over descriptors + +use crate::descriptor::{TapTreeIter, Tr}; +use crate::miniscript::context::{BareCtx, Legacy, Segwitv0, Tap}; +use crate::{miniscript, Miniscript, MiniscriptKey}; + +/// Iterator over all the keys in a descriptor. +pub struct PkIter<'desc, Pk: MiniscriptKey> { + single_key: Option, + taptree_iter: Option>, + ms_iter_bare: Option>, + ms_iter_legacy: Option>, + ms_iter_segwit: Option>, + ms_iter_taproot: Option>, + sorted_multi: Option>, +} + +impl<'desc, Pk: MiniscriptKey> PkIter<'desc, Pk> { + pub(super) fn from_key(pk: Pk) -> Self { + Self { + single_key: Some(pk), + taptree_iter: None, + ms_iter_bare: None, + ms_iter_legacy: None, + ms_iter_segwit: None, + ms_iter_taproot: None, + sorted_multi: None, + } + } + + pub(super) fn from_miniscript_bare(ms: &'desc Miniscript) -> Self { + Self { + single_key: None, + taptree_iter: None, + ms_iter_bare: Some(ms.iter_pk()), + ms_iter_legacy: None, + ms_iter_segwit: None, + ms_iter_taproot: None, + sorted_multi: None, + } + } + + pub(super) fn from_miniscript_legacy(ms: &'desc Miniscript) -> Self { + Self { + single_key: None, + taptree_iter: None, + ms_iter_bare: None, + ms_iter_legacy: Some(ms.iter_pk()), + ms_iter_segwit: None, + ms_iter_taproot: None, + sorted_multi: None, + } + } + + pub(super) fn from_miniscript_segwit(ms: &'desc Miniscript) -> Self { + Self { + single_key: None, + taptree_iter: None, + ms_iter_bare: None, + ms_iter_legacy: None, + ms_iter_segwit: Some(ms.iter_pk()), + ms_iter_taproot: None, + sorted_multi: None, + } + } + + pub(super) fn from_sortedmulti(sm: &'desc [Pk]) -> Self { + Self { + single_key: None, + taptree_iter: None, + ms_iter_bare: None, + ms_iter_legacy: None, + ms_iter_segwit: None, + ms_iter_taproot: None, + sorted_multi: Some(sm.iter()), + } + } + + pub(super) fn from_tr(tr: &'desc Tr) -> Self { + Self { + single_key: Some(tr.internal_key().clone()), + taptree_iter: Some(tr.leaves()), + ms_iter_bare: None, + ms_iter_legacy: None, + ms_iter_segwit: None, + ms_iter_taproot: None, + sorted_multi: None, + } + } +} + +impl<'desc, Pk: MiniscriptKey> Iterator for PkIter<'desc, Pk> { + type Item = Pk; + + #[rustfmt::skip] // the tower of .or_elses looks good as is + fn next(&mut self) -> Option { + // If there is a single key, return it first. (This will be the case + // for all single-key-only iterators but also for Taproot, where the + // single key is the root key.) + if let Some(k) = self.single_key.take() { + return Some(k.clone()); + } + + // Then attempt to yield something from the Taptree iterator. + loop { + if let Some(item) = self.ms_iter_taproot.as_mut().and_then(Iterator::next) { + return Some(item); + } + if let Some(iter) = self.taptree_iter.as_mut().and_then(Iterator::next) { + self.ms_iter_taproot = Some(iter.miniscript().iter_pk()); + } else { + break; + } + } + + // Finally run through the train of other iterators. + self.ms_iter_bare.as_mut().and_then(Iterator::next).or_else( + || self.ms_iter_legacy.as_mut().and_then(Iterator::next).or_else( + || self.ms_iter_segwit.as_mut().and_then(Iterator::next).or_else( + || self.sorted_multi.as_mut().and_then(Iterator::next).cloned() + ) + ) + ) + } +} diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs index 4d7afa4b2..87b70b98f 100644 --- a/src/descriptor/mod.rs +++ b/src/descriptor/mod.rs @@ -32,6 +32,7 @@ use crate::{ }; mod bare; +mod iter; mod segwitv0; mod sh; mod sortedmulti; @@ -39,6 +40,7 @@ mod tr; // Descriptor Exports pub use self::bare::{Bare, Pkh}; +pub use self::iter::PkIter; pub use self::segwitv0::{Wpkh, Wsh, WshInner}; pub use self::sh::{Sh, ShInner}; pub use self::sortedmulti::SortedMultiVec; @@ -241,6 +243,29 @@ impl Descriptor { Ok(Descriptor::Tr(Tr::new(key, script)?)) } + /// An iterator over all the keys referenced in the descriptor. + pub fn iter_pk(&self) -> PkIter<'_, Pk> { + match *self { + Descriptor::Bare(ref bare) => PkIter::from_miniscript_bare(bare.as_inner()), + Descriptor::Pkh(ref pk) => PkIter::from_key(pk.as_inner().clone()), + Descriptor::Wpkh(ref pk) => PkIter::from_key(pk.as_inner().clone()), + Descriptor::Sh(ref sh) => match *sh.as_inner() { + ShInner::Wsh(ref wsh) => match wsh.as_inner() { + WshInner::SortedMulti(ref sorted) => PkIter::from_sortedmulti(sorted.pks()), + WshInner::Ms(ref ms) => PkIter::from_miniscript_segwit(ms), + }, + ShInner::Wpkh(ref pk) => PkIter::from_key(pk.as_inner().clone()), + ShInner::SortedMulti(ref sorted) => PkIter::from_sortedmulti(sorted.pks()), + ShInner::Ms(ref ms) => PkIter::from_miniscript_legacy(ms), + }, + Descriptor::Wsh(ref wsh) => match wsh.as_inner() { + WshInner::SortedMulti(ref sorted) => PkIter::from_sortedmulti(sorted.pks()), + WshInner::Ms(ref ms) => PkIter::from_miniscript_segwit(ms), + }, + Descriptor::Tr(ref tr) => PkIter::from_tr(tr), + } + } + /// For a Taproot descriptor, returns the internal key. pub fn internal_key(&self) -> Option<&Pk> { if let Descriptor::Tr(ref tr) = self { From aa42cd390a700384902fb932c63c308ee0486272 Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Fri, 23 May 2025 23:13:05 +0000 Subject: [PATCH 2/2] descriptor: add unit tests for iter_pk This one -was- AI-generated, though I reviewed it. --- src/descriptor/mod.rs | 169 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 169 insertions(+) diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs index 87b70b98f..c2ee02b71 100644 --- a/src/descriptor/mod.rs +++ b/src/descriptor/mod.rs @@ -2262,4 +2262,173 @@ pk(03f28773c2d975288bc7d1d205c3748651b075fbc6610e58cddeeddf8f19405aa8))"; assert_eq!(xonly_pk_descriptor.to_string(), xonly_converted_descriptor.to_string()); } + + #[test] + fn test_iter_pk() { + // Test Bare descriptor + let bare_desc = Descriptor::::from_str( + "pk(020000000000000000000000000000000000000000000000000000000000000002)", + ) + .unwrap(); + let keys: Vec<_> = bare_desc.iter_pk().collect(); + assert_eq!(keys.len(), 1); + assert_eq!( + keys[0].to_string(), + "020000000000000000000000000000000000000000000000000000000000000002" + ); + + // Test Pkh descriptor + let pkh_desc = Descriptor::::from_str( + "pkh(020000000000000000000000000000000000000000000000000000000000000002)", + ) + .unwrap(); + let keys: Vec<_> = pkh_desc.iter_pk().collect(); + assert_eq!(keys.len(), 1); + assert_eq!( + keys[0].to_string(), + "020000000000000000000000000000000000000000000000000000000000000002" + ); + + // Test Wpkh descriptor + let wpkh_desc = Descriptor::::from_str( + "wpkh(020000000000000000000000000000000000000000000000000000000000000002)", + ) + .unwrap(); + let keys: Vec<_> = wpkh_desc.iter_pk().collect(); + assert_eq!(keys.len(), 1); + assert_eq!( + keys[0].to_string(), + "020000000000000000000000000000000000000000000000000000000000000002" + ); + + // Test Sh descriptor with a miniscript + let sh_desc = Descriptor::::from_str( + "sh(or_d(pk(021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b),pk(0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6)))" + ).unwrap(); + let keys: Vec<_> = sh_desc.iter_pk().collect(); + assert_eq!(keys.len(), 2); + assert_eq!( + keys[0].to_string(), + "021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b" + ); + assert_eq!( + keys[1].to_string(), + "0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6" + ); + + // Test Sh-Wpkh descriptor + let shwpkh_desc = Descriptor::::from_str( + "sh(wpkh(020000000000000000000000000000000000000000000000000000000000000002))", + ) + .unwrap(); + let keys: Vec<_> = shwpkh_desc.iter_pk().collect(); + assert_eq!(keys.len(), 1); + assert_eq!( + keys[0].to_string(), + "020000000000000000000000000000000000000000000000000000000000000002" + ); + + // Test Sh-Wsh descriptor + let shwsh_desc = Descriptor::::from_str( + "sh(wsh(or_d(pk(021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b),pk(0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6))))" + ).unwrap(); + let keys: Vec<_> = shwsh_desc.iter_pk().collect(); + assert_eq!(keys.len(), 2); + assert_eq!( + keys[0].to_string(), + "021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b" + ); + assert_eq!( + keys[1].to_string(), + "0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6" + ); + + // Test Wsh descriptor + let wsh_desc = Descriptor::::from_str( + "wsh(or_d(pk(021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b),pk(0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6)))" + ).unwrap(); + let keys: Vec<_> = wsh_desc.iter_pk().collect(); + assert_eq!(keys.len(), 2); + assert_eq!( + keys[0].to_string(), + "021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b" + ); + assert_eq!( + keys[1].to_string(), + "0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6" + ); + + // Test SortedMulti descriptors + let sortedmulti_desc = Descriptor::::from_str( + "sh(sortedmulti(2,021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b,0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6,03fff97bd5755eeea420453a14355235d382f6472f8568a18b2f057a1460297556))" + ).unwrap(); + let keys: Vec<_> = sortedmulti_desc.iter_pk().collect(); + assert_eq!(keys.len(), 3); + // Keys are sorted in the output + assert!(keys.iter().any(|k| k.to_string() + == "021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b")); + assert!(keys.iter().any(|k| k.to_string() + == "0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6")); + assert!(keys.iter().any(|k| k.to_string() + == "03fff97bd5755eeea420453a14355235d382f6472f8568a18b2f057a1460297556")); + + // Test Taproot descriptor with only key path + let tr_key_only_desc = Descriptor::::from_str( + "tr(020000000000000000000000000000000000000000000000000000000000000002)", + ) + .unwrap(); + let keys: Vec<_> = tr_key_only_desc.iter_pk().collect(); + assert_eq!(keys.len(), 1); + assert_eq!( + keys[0].to_string(), + "020000000000000000000000000000000000000000000000000000000000000002" + ); + + // Test Taproot descriptor with script path + // The internal key should be yielded first + let internal_key = "020000000000000000000000000000000000000000000000000000000000000002"; + let script_key1 = "021d4ea7132d4e1a362ee5efd8d0b59dd4d1fe8906eefa7dd812b05a46b73d829b"; + let script_key2 = "0302c8bbbb393f32c843149ce36d56405595aaabab2d0e1f4ca5f9de67dd7419f6"; + + let tr_with_script_desc = Descriptor::::from_str(&format!( + "tr({},{{pk({}),pk({})}})", + internal_key, script_key1, script_key2, + )) + .unwrap(); + + let keys: Vec<_> = tr_with_script_desc.iter_pk().collect(); + assert_eq!(keys.len(), 3); + + // Verify internal key is first + assert_eq!(keys[0].to_string(), internal_key); + + // Verify other keys are present (order after internal key is not guaranteed) + assert!(keys[1..].iter().any(|k| k.to_string() == script_key1)); + assert!(keys[1..].iter().any(|k| k.to_string() == script_key2)); + + // Test Taproot descriptor with complex script tree + let tr_complex_desc = Descriptor::::from_str(&format!( + "tr({},{{pk({}),{{pk({}),or_d(pk({}),pk({}))}}}})", + internal_key, + script_key1, + script_key2, + "03fff97bd5755eeea420453a14355235d382f6472f8568a18b2f057a1460297556", + "0250863ad64a87ae8a2fe83c1af1a8403cb53f53e486d8511dad8a04887e5b2352" + )) + .unwrap(); + + let keys: Vec<_> = tr_complex_desc.iter_pk().collect(); + assert_eq!(keys.len(), 5); + + // Verify internal key is first + assert_eq!(keys[0].to_string(), internal_key); + + // Verify all other keys are present + assert!(keys[1..].iter().any(|k| k.to_string() == script_key1)); + assert!(keys[1..].iter().any(|k| k.to_string() == script_key2)); + assert!(keys[1..].iter().any(|k| k.to_string() + == "03fff97bd5755eeea420453a14355235d382f6472f8568a18b2f057a1460297556")); + assert!(keys[1..].iter().any(|k| k.to_string() + == "0250863ad64a87ae8a2fe83c1af1a8403cb53f53e486d8511dad8a04887e5b2352")); + } }