Skip to content

Add Descriptor::iter_pk #823

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

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add Descriptor::iter_pk #823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
523654e
descriptor: add Descriptor::iter_pk
apoelstra May 23, 2025
aa42cd3
descriptor: add unit tests for iter_pk
apoelstra May 23, 2025
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
127 changes: 127 additions & 0 deletions src/descriptor/iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is fine. But Ideally, we want this to be a sum type instead of a product type? But then we get whole loads of matches everywhere?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The main reason to use a product type was that Taproot descriptors have both single keys and (optionally) a taptree iterator.

And in PkIter::next I have a construction where the taptree iterator yields new Miniscript pkiters.

I can try to refactor this to use a sum type but I think it would result in more repeated/redundant code.

single_key: Option<Pk>,
taptree_iter: Option<TapTreeIter<'desc, Pk>>,
ms_iter_bare: Option<miniscript::iter::PkIter<'desc, Pk, BareCtx>>,
ms_iter_legacy: Option<miniscript::iter::PkIter<'desc, Pk, Legacy>>,
ms_iter_segwit: Option<miniscript::iter::PkIter<'desc, Pk, Segwitv0>>,
ms_iter_taproot: Option<miniscript::iter::PkIter<'desc, Pk, Tap>>,
sorted_multi: Option<core::slice::Iter<'desc, Pk>>,
}

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<Pk, BareCtx>) -> 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<Pk, Legacy>) -> 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<Pk, Segwitv0>) -> 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<Pk>) -> 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<Self::Item> {
// 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()
)
)
)
}
}
194 changes: 194 additions & 0 deletions src/descriptor/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,15 @@ use crate::{
};

mod bare;
mod iter;
mod segwitv0;
mod sh;
mod sortedmulti;
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;
Expand Down Expand Up @@ -241,6 +243,29 @@ impl<Pk: MiniscriptKey> Descriptor<Pk> {
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 {
Expand Down Expand Up @@ -2237,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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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::<PublicKey>::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"));
}
}
Loading