diff --git a/src/miniscript/musig_key.rs b/src/miniscript/musig_key.rs index 357fdfa97..8b2290270 100644 --- a/src/miniscript/musig_key.rs +++ b/src/miniscript/musig_key.rs @@ -40,6 +40,35 @@ impl FromStr for KeyExpr { } } +#[derive(Debug, Clone)] +/// Iterator for keyexpr +pub struct KeyExprIter<'a, Pk: MiniscriptKey> { + stack: Vec<&'a KeyExpr>, +} + +impl<'a, Pk> Iterator for KeyExprIter<'a, Pk> +where + Pk: MiniscriptKey + 'a, +{ + type Item = &'a Pk; + + fn next(&mut self) -> Option { + while !self.stack.is_empty() { + let last = self.stack.pop().expect("Size checked above"); + match &*last { + KeyExpr::MuSig(key_vec) => { + // push the elements in reverse order + for key in key_vec.iter().rev() { + self.stack.push(key) + } + } + KeyExpr::SingleKey(ref pk) => return Some(pk), + } + } + None + } +} + impl fmt::Debug for KeyExpr { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match *self { @@ -79,6 +108,14 @@ impl fmt::Display for KeyExpr { } } } + +impl KeyExpr { + /// Iterate over all keys + pub fn iter(&self) -> KeyExprIter { + KeyExprIter { stack: vec![self] } + } +} + #[cfg(test)] mod tests { use super::*; @@ -97,4 +134,33 @@ mod tests { test_one("A"); test_one("musig(,,)"); } + + #[test] + fn test_iterator() { + let pk = KeyExpr::::from_str("musig(A,B,musig(C,musig(D,E)))").unwrap(); + let mut my_iter = pk.iter(); + assert_eq!(my_iter.next(), Some(&String::from("A"))); + assert_eq!(my_iter.next(), Some(&String::from("B"))); + assert_eq!(my_iter.next(), Some(&String::from("C"))); + assert_eq!(my_iter.next(), Some(&String::from("D"))); + assert_eq!(my_iter.next(), Some(&String::from("E"))); + assert_eq!(my_iter.next(), None); + } + + fn test_helper(musig_key: &str, comma_separated_key: &str) { + let pk = KeyExpr::::from_str(musig_key).unwrap(); + let var: Vec<&str> = comma_separated_key.split(",").collect(); + let key_names: Vec<&String> = pk.iter().collect(); + for (key1, key2) in key_names.iter().zip(var.iter()) { + assert_eq!(key1, key2); + } + } + + #[test] + fn test_iterator_multi() { + test_helper("musig(A)", "A"); + test_helper("A", "A"); + test_helper("musig(,,)", ""); + test_helper("musig(musig(A,B),musig(musig(C)))", "A,B,C"); + } }