Add Iterator for KeyExpr

Author: kanishk779

src/miniscript/musig_key.rs

@@ -40,6 +40,35 @@ impl<Pk: MiniscriptKey + FromStr> FromStr for KeyExpr<Pk> {
     }
 }
 
+#[derive(Debug, Clone)]
+/// Iterator for keyexpr
+pub struct KeyExprIter<'a, Pk: MiniscriptKey> {
+    stack: Vec<&'a KeyExpr<Pk>>,
+}
+
+impl<'a, Pk> Iterator for KeyExprIter<'a, Pk>
+where
+    Pk: MiniscriptKey + 'a,
+{
+    type Item = &'a Pk;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        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<Pk: MiniscriptKey> fmt::Debug for KeyExpr<Pk> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
@@ -79,6 +108,14 @@ impl<Pk: MiniscriptKey> fmt::Display for KeyExpr<Pk> {
         }
     }
 }
+
+impl<Pk: MiniscriptKey> KeyExpr<Pk> {
+    /// Iterate over all keys
+    pub fn iter(&self) -> KeyExprIter<Pk> {
+        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::<String>::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::<String>::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");
+    }
 }