descriptor: add Descriptor::iter_pk

This commit was -not- AI-generated, though probably it should've been.

Author: apoelstra

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<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()
+                )
+            )
+        )
+    }
+}

src/descriptor/mod.rs

@@ -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;
@@ -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 {