Policy to single-leaf TapTree compilation done

Author: SarcasticNastik

src/policy/compiler.rs

@@ -22,6 +22,7 @@ use std::convert::From;
 use std::marker::PhantomData;
 use std::{cmp, error, f64, fmt, mem};
 
+use miniscript::context::SigType;
 use miniscript::limits::MAX_PUBKEYS_PER_MULTISIG;
 use miniscript::types::{self, ErrorKind, ExtData, Property, Type};
 use miniscript::ScriptContext;
@@ -994,7 +995,9 @@ where
                 })
                 .collect();
 
-            if key_vec.len() == subs.len() && subs.len() <= MAX_PUBKEYS_PER_MULTISIG {
+            if Ctx::sig_type() == SigType::Schnorr {
+                insert_wrap!(AstElemExt::terminal(Terminal::MultiA(k, key_vec)));
+            } else if key_vec.len() == subs.len() && subs.len() <= MAX_PUBKEYS_PER_MULTISIG {
                 insert_wrap!(AstElemExt::terminal(Terminal::Multi(k, key_vec)));
             }
             // Not a threshold, it's always more optimal to translate it to and()s as we save the
@@ -1168,7 +1171,7 @@ mod tests {
 
     use miniscript::{satisfy, Legacy, Segwitv0};
     use policy::Liftable;
-    use script_num_size;
+    use {script_num_size, Tap};
 
     type SPolicy = Concrete<String>;
     type BPolicy = Concrete<bitcoin::PublicKey>;
@@ -1554,6 +1557,17 @@ mod tests {
             ))
         );
     }
+
+    #[test]
+    fn compile_tr_thresh() {
+        for k in 1..4 {
+            let small_thresh: Concrete<String> =
+                policy_str!("{}", &format!("thresh({},pk(B),pk(C),pk(D))", k));
+            let small_thresh_ms: Miniscript<String, Tap> = small_thresh.compile().unwrap();
+            let small_thresh_ms_expected: Miniscript<String, Tap> = ms_str!("multi_a({},B,C,D)", k);
+            assert_eq!(small_thresh_ms, small_thresh_ms_expected);
+        }
+    }
 }
 
 #[cfg(all(test, feature = "unstable"))]

src/policy/concrete.rs

@@ -18,9 +18,13 @@
 use bitcoin::hashes::hex::FromHex;
 use bitcoin::hashes::{hash160, ripemd160, sha256, sha256d};
 use std::collections::HashSet;
+#[cfg(feature = "compiler")]
+use std::sync::Arc;
 use std::{error, fmt, str};
 
 use super::ENTAILMENT_MAX_TERMINALS;
+#[cfg(feature = "compiler")]
+use descriptor::TapTree;
 use errstr;
 use expression::{self, FromTree};
 use miniscript::limits::{HEIGHT_TIME_THRESHOLD, SEQUENCE_LOCKTIME_TYPE_FLAG};
@@ -32,8 +36,13 @@ use policy::compiler;
 #[cfg(feature = "compiler")]
 use policy::compiler::CompilerError;
 #[cfg(feature = "compiler")]
+use Descriptor;
+#[cfg(feature = "compiler")]
 use Miniscript;
+#[cfg(feature = "compiler")]
+use Tap;
 use {Error, ForEach, ForEachKey, MiniscriptKey};
+
 /// Concrete policy which corresponds directly to a Miniscript structure,
 /// and whose disjunctions are annotated with satisfaction probabilities
 /// to assist the compiler
@@ -128,6 +137,33 @@ impl fmt::Display for PolicyError {
 }
 
 impl<Pk: MiniscriptKey> Policy<Pk> {
+    /// Single-Node compilation
+    #[cfg(feature = "compiler")]
+    fn compile_huffman_taptree(policy: &Self) -> Result<TapTree<Pk>, Error> {
+        let compilation = policy.compile::<Tap>().unwrap();
+        Ok(TapTree::Leaf(Arc::new(compilation)))
+    }
+    /// Extract the Taproot internal_key from policy tree.
+    #[cfg(feature = "compiler")]
+    fn extract_key(policy: &Self, unspendable_key: Option<Pk>) -> Result<(Pk, &Policy<Pk>), Error> {
+        match unspendable_key {
+            Some(key) => Ok((key, policy)),
+            None => Err(errstr("No internal key found")),
+        }
+    }
+
+    /// Compile the [`Tr`] descriptor into optimized [`TapTree`] implementation
+    #[cfg(feature = "compiler")]
+    pub fn compile_tr(&self, unspendable_key: Option<Pk>) -> Result<Descriptor<Pk>, Error> {
+        let (internal_key, policy) = Self::extract_key(self, unspendable_key).unwrap();
+        let tree = Descriptor::new_tr(
+            internal_key,
+            Some(Self::compile_huffman_taptree(policy).unwrap()),
+        )
+        .unwrap();
+        Ok(tree)
+    }
+
     /// Compile the descriptor into an optimized `Miniscript` representation
     #[cfg(feature = "compiler")]
     pub fn compile<Ctx: ScriptContext>(&self) -> Result<Miniscript<Pk, Ctx>, CompilerError> {

src/policy/mod.rs

@@ -233,8 +233,14 @@ mod tests {
         Concrete, Liftable, Semantic,
     };
     use bitcoin;
+    #[cfg(feature = "compiler")]
+    use descriptor::TapTree;
     use std::str::FromStr;
+    #[cfg(feature = "compiler")]
+    use std::sync::Arc;
     use DummyKey;
+    #[cfg(feature = "compiler")]
+    use {Descriptor, Tap};
 
     type ConcretePol = Concrete<DummyKey>;
     type SemanticPol = Semantic<DummyKey>;
@@ -364,4 +370,20 @@ mod tests {
             ms_str.lift().unwrap()
         );
     }
+
+    #[test]
+    #[cfg(feature = "compiler")]
+    fn single_leaf_tr_compile() {
+        for k in 1..5 {
+            let unspendable_key: String = "z".to_string();
+            let policy: Concrete<String> = policy_str!("thresh({},pk(A),pk(B),pk(C),pk(D))", k);
+            let descriptor = policy.compile_tr(Some(unspendable_key.clone())).unwrap();
+
+            let ms_compilation: Miniscript<String, Tap> = ms_str!("multi_a({},A,B,C,D)", k);
+            let tree: TapTree<String> = TapTree::Leaf(Arc::new(ms_compilation));
+            let expected_descriptor = Descriptor::new_tr(unspendable_key, Some(tree)).unwrap();
+
+            assert_eq!(descriptor, expected_descriptor);
+        }
+    }
 }