From 009186f361ae5f4553a18eaea85774a0ec7a2212 Mon Sep 17 00:00:00 2001
From: kanishk779 <kashyapamankumar54@gmail.com>
Date: Fri, 1 Jul 2022 23:46:39 +0530
Subject: [PATCH] Add initial support for KeyExpr

---
 src/miniscript/mod.rs       |   1 +
 src/miniscript/musig_key.rs | 102 ++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 src/miniscript/musig_key.rs

diff --git a/src/miniscript/mod.rs b/src/miniscript/mod.rs
index 9d50d63d1..f835364ce 100644
--- a/src/miniscript/mod.rs
+++ b/src/miniscript/mod.rs
@@ -41,6 +41,7 @@ pub mod hash256;
 pub mod iter;
 pub mod lex;
 pub mod limits;
+pub mod musig_key;
 pub mod satisfy;
 pub mod types;
 
diff --git a/src/miniscript/musig_key.rs b/src/miniscript/musig_key.rs
new file mode 100644
index 000000000..5cf2c576b
--- /dev/null
+++ b/src/miniscript/musig_key.rs
@@ -0,0 +1,102 @@
+//! Support for multi-signature keys
+use std;
+use std::str::FromStr;
+use std::{error, fmt};
+
+use crate::expression::Tree;
+use crate::MiniscriptKey;
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq, Hash)]
+/// Enum for representing keys in miniscript
+pub enum KeyExpr<Pk: MiniscriptKey> {
+    /// Single-key (e.g pk(a), here 'a' is a single key)
+    SingleKey(Pk),
+
+    /// Collection of keys in used for multi-signature
+    MuSig(Vec<KeyExpr<Pk>>),
+}
+
+#[derive(Debug, PartialEq)]
+/// Errors related to KeyExpr
+pub enum KeyExprError {
+    /// Single key is empty string
+    SingleKeyEmptyError,
+
+    /// Parsing error for single key
+    SingleKeyParseError,
+}
+
+impl fmt::Display for KeyExprError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            KeyExprError::SingleKeyEmptyError => f.write_str("str for single key was empty"),
+            KeyExprError::SingleKeyParseError => f.write_str("not able to parse the single key"),
+        }
+    }
+}
+
+impl error::Error for KeyExprError {}
+
+impl<Pk: MiniscriptKey + FromStr> FromStr for KeyExpr<Pk> {
+    type Err = KeyExprError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let (key_tree, _) = Tree::from_slice(s).unwrap();
+        fn expr_from_tree<Pk: MiniscriptKey + FromStr>(
+            tree: &Tree,
+        ) -> Result<KeyExpr<Pk>, KeyExprError> {
+            if tree.name == "musig" {
+                let mut key_expr_vect = vec![];
+                for sub_tree in tree.args.iter() {
+                    let temp_res = expr_from_tree(sub_tree)?;
+                    key_expr_vect.push(temp_res);
+                }
+                Ok(KeyExpr::MuSig(key_expr_vect))
+            } else {
+                if tree.name != "" {
+                    let single_key = match Pk::from_str(tree.name) {
+                        Ok(x) => x,
+                        Err(_) => {
+                            return Err(KeyExprError::SingleKeyParseError);
+                        }
+                    };
+                    Ok(KeyExpr::SingleKey(single_key))
+                } else {
+                    Err(KeyExprError::SingleKeyEmptyError)
+                }
+            }
+        }
+        expr_from_tree(&key_tree)
+    }
+}
+
+impl<Pk: MiniscriptKey> fmt::Display for KeyExpr<Pk> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            KeyExpr::SingleKey(ref pk) => write!(f, "{}", pk),
+            KeyExpr::MuSig(ref my_vec) => {
+                write!(f, "musig(")?;
+                let len = my_vec.len();
+                for (index, k) in my_vec.iter().enumerate() {
+                    if index == len - 1 {
+                        write!(f, "{}", k)?;
+                    } else {
+                        write!(f, "{},", k)?;
+                    }
+                }
+                f.write_str(")")
+            }
+        }
+    }
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_one() {
+        let dummy_key = "musig(A,B,musig(C,musig(D,E)))";
+        let pk = KeyExpr::<String>::from_str(dummy_key).unwrap();
+        println!("{}", pk);
+        assert_eq!(dummy_key, format!("{}", pk))
+    }
+}