Policy to single-leaf TapTree compilation done

SarcasticNastik · SarcasticNastik · commit 26fc574434e9 · 2022-05-11T14:36:50.000+05:30
A new method `compile_tr` is introduced in policy for creating a Tr descriptor with TapTree containing single script child-node. An internal-key extraction method from the policy has also been implemented for the same.
diff --git a/src/policy/compiler.rs b/src/policy/compiler.rs
@@ -24,6 +24,7 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 use std::{cmp, error, f64, fmt, hash, mem};
 
+use crate::miniscript::context::SigType;
 use crate::miniscript::limits::MAX_PUBKEYS_PER_MULTISIG;
 use crate::miniscript::types::{self, ErrorKind, ExtData, Property, Type};
 use crate::miniscript::ScriptContext;
@@ -987,18 +988,23 @@ where
                 })
                 .collect();
 
-            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
-            // resulting threshold check (N EQUAL) in any case.
-            else if k == subs.len() {
-                let mut policy = subs.first().expect("No sub policy in thresh() ?").clone();
-                for sub in &subs[1..] {
-                    policy = Concrete::And(vec![sub.clone(), policy]);
+            match Ctx::sig_type() {
+                SigType::Schnorr if key_vec.len() == subs.len() => {
+                    insert_wrap!(AstElemExt::terminal(Terminal::MultiA(k, key_vec)))
+                }
+                SigType::Ecdsa
+                    if key_vec.len() == subs.len() && subs.len() <= MAX_PUBKEYS_PER_MULTISIG =>
+                {
+                    insert_wrap!(AstElemExt::terminal(Terminal::Multi(k, key_vec)))
                 }
+                _ if k == subs.len() => {
+                    let mut it = subs.iter();
+                    let mut policy = it.next().expect("No sub policy in thresh() ?").clone();
+                    policy = it.fold(policy, |acc, pol| Concrete::And(vec![acc, pol.clone()]));
 
-                ret = best_compilations(policy_cache, &policy, sat_prob, dissat_prob)?;
+                    ret = best_compilations(policy_cache, &policy, sat_prob, dissat_prob)?;
+                }
+                _ => {}
             }
 
             // FIXME: Should we also optimize thresh(1, subs) ?
@@ -1549,6 +1555,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"))]
diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs
@@ -20,20 +20,19 @@ use std::{error, fmt, str};
 
 use bitcoin::hashes::hex::FromHex;
 use bitcoin::hashes::{hash160, ripemd160, sha256, sha256d};
+#[cfg(feature = "compiler")]
+use {
+    crate::descriptor::TapTree, crate::miniscript::ScriptContext, crate::policy::compiler,
+    crate::policy::compiler::CompilerError, crate::Descriptor, crate::Miniscript, crate::Tap,
+    std::sync::Arc,
+};
 
 use super::ENTAILMENT_MAX_TERMINALS;
 use crate::expression::{self, FromTree};
 use crate::miniscript::limits::{HEIGHT_TIME_THRESHOLD, SEQUENCE_LOCKTIME_TYPE_FLAG};
 use crate::miniscript::types::extra_props::TimeLockInfo;
-#[cfg(feature = "compiler")]
-use crate::miniscript::ScriptContext;
-#[cfg(feature = "compiler")]
-use crate::policy::compiler;
-#[cfg(feature = "compiler")]
-use crate::policy::compiler::CompilerError;
-#[cfg(feature = "compiler")]
-use crate::Miniscript;
 use crate::{errstr, 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 +127,21 @@ impl fmt::Display for PolicyError {
 }
 
 impl<Pk: MiniscriptKey> Policy<Pk> {
+    /// Single-Node compilation
+    #[cfg(feature = "compiler")]
+    fn compile_leaf_taptree(&self) -> Result<TapTree<Pk>, Error> {
+        let compilation = self.compile::<Tap>().unwrap();
+        Ok(TapTree::Leaf(Arc::new(compilation)))
+    }
+
+    /// Compile a [`Policy`] into a single-leaf [`TapTree`]
+    #[cfg(feature = "compiler")]
+    pub fn compile_tr(&self, unspendable_key: Option<Pk>) -> Result<Descriptor<Pk>, Error> {
+        let internal_key = unspendable_key.ok_or(errstr("No internal key found"))?;
+        let tree = Descriptor::new_tr(internal_key, Some(self.compile_leaf_taptree()?))?;
+        Ok(tree)
+    }
+
     /// Compile the descriptor into an optimized `Miniscript` representation
     #[cfg(feature = "compiler")]
     pub fn compile<Ctx: ScriptContext>(&self) -> Result<Miniscript<Pk, Ctx>, CompilerError> {
diff --git a/src/policy/mod.rs b/src/policy/mod.rs
@@ -225,13 +225,17 @@ impl<Pk: MiniscriptKey> Liftable<Pk> for Concrete<Pk> {
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;
+    #[cfg(feature = "compiler")]
+    use std::sync::Arc;
 
     use bitcoin;
 
     use super::super::miniscript::context::Segwitv0;
     use super::super::miniscript::Miniscript;
     use super::{Concrete, Liftable, Semantic};
     use crate::DummyKey;
+    #[cfg(feature = "compiler")]
+    use crate::{descriptor::TapTree, Descriptor, Tap};
 
     type ConcretePol = Concrete<DummyKey>;
     type SemanticPol = Semantic<DummyKey>;
@@ -361,4 +365,18 @@ mod tests {
             ms_str.lift().unwrap()
         );
     }
+
+    #[test]
+    #[cfg(feature = "compiler")]
+    fn single_leaf_tr_compile() {
+        let unspendable_key: String = "z".to_string();
+        let policy: Concrete<String> = policy_str!("thresh(2,pk(A),pk(B),pk(C),pk(D))");
+        let descriptor = policy.compile_tr(Some(unspendable_key.clone())).unwrap();
+
+        let ms_compilation: Miniscript<String, Tap> = ms_str!("multi_a(2,A,B,C,D)");
+        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);
+    }
 }
