trivial P2Tr compiler done

SarcasticNastik · SarcasticNastik · commit b10cb7a6e8c8 · 2022-01-16T21:50:29.000+05:30
Taproot internal-key extraction from policy done
diff --git a/src/policy/compiler.rs b/src/policy/compiler.rs
@@ -38,7 +38,7 @@ type PolicyCache<Pk, Ctx> =
 
 ///Ordered f64 for comparison
 #[derive(Copy, Clone, PartialEq, PartialOrd, Debug)]
-struct OrdF64(f64);
+pub struct OrdF64(pub f64);
 
 impl Eq for OrdF64 {}
 impl Ord for OrdF64 {
diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs
@@ -18,10 +18,10 @@
 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};
 
+#[cfg(feature = "compiler")]
+use super::Liftable;
 use super::ENTAILMENT_MAX_TERMINALS;
 #[cfg(feature = "compiler")]
 use descriptor::TapTree;
@@ -32,9 +32,15 @@ use miniscript::types::extra_props::TimeLockInfo;
 #[cfg(feature = "compiler")]
 use miniscript::ScriptContext;
 #[cfg(feature = "compiler")]
-use policy::compiler;
+use policy::compiler::{CompilerError, OrdF64};
+#[cfg(feature = "compiler")]
+use policy::{compiler, Semantic};
+#[cfg(feature = "compiler")]
+use std::cmp::Reverse;
 #[cfg(feature = "compiler")]
-use policy::compiler::CompilerError;
+use std::collections::BinaryHeap;
+#[cfg(feature = "compiler")]
+use std::sync::Arc;
 #[cfg(feature = "compiler")]
 use Descriptor;
 #[cfg(feature = "compiler")]
@@ -137,30 +143,109 @@ impl fmt::Display for PolicyError {
 }
 
 impl<Pk: MiniscriptKey> Policy<Pk> {
-    /// Single-Node compilation
+    /// Create a Huffman Tree from compiled [Miniscript] nodes
+    #[cfg(feature = "compiler")]
+    fn with_huffman_tree<T>(
+        ms: Vec<(OrdF64, Miniscript<Pk, Tap>)>,
+        f: T,
+    ) -> Result<TapTree<Pk>, Error>
+    where
+        T: Fn(OrdF64) -> OrdF64,
+    {
+        // Pattern match terminal Or/ Terminal (with equal odds)
+        let mut node_weights = BinaryHeap::<(Reverse<OrdF64>, Arc<TapTree<Pk>>)>::new();
+        for (prob, script) in ms {
+            node_weights.push((Reverse(f(prob)), Arc::from(TapTree::Leaf(Arc::new(script)))));
+        }
+        if node_weights.is_empty() {
+            return Err(errstr("Empty Miniscript compilation"));
+        }
+        while node_weights.len() > 1 {
+            let (p1, s1) = node_weights.pop().expect("len must atleast be two");
+            let (p2, s2) = node_weights.pop().expect("len must atleast be two");
+
+            let p = (p1.0).0 + (p2.0).0;
+            node_weights.push((Reverse(OrdF64(p)), Arc::from(TapTree::Tree(s1, s2))));
+        }
+
+        debug_assert!(node_weights.len() == 1);
+        let node = node_weights
+            .pop()
+            .expect("huffman tree algorithm is broken")
+            .1;
+        Ok((*node).clone())
+    }
+
+    /// Flatten the [`Policy`] tree structure into a Vector with corresponding leaf probability
+    // TODO: 1. Can try to push the maximum of scaling factors and accordingly update later for
+    // TODO: 1. integer metric. (Accordingly change metrics everywhere)
+    #[cfg(feature = "compiler")]
+    fn flatten_policy(&self, prob: f64) -> Vec<(f64, Policy<Pk>)> {
+        match *self {
+            Policy::Or(ref subs) => {
+                let total_odds: usize = subs.iter().map(|(ref k, _)| k).sum();
+                subs.iter()
+                    .map(|(k, ref policy)| {
+                        policy.flatten_policy(prob * *k as f64 / total_odds as f64)
+                    })
+                    .flatten()
+                    .collect::<Vec<_>>()
+            }
+            Policy::Threshold(k, ref subs) if k == 1 => {
+                let total_odds = subs.len();
+                subs.iter()
+                    .map(|policy| policy.flatten_policy(prob / total_odds as f64))
+                    .flatten()
+                    .collect::<Vec<_>>()
+            }
+            ref x => vec![(prob, x.clone())],
+        }
+    }
+
+    /// Compile [`Policy::Or`] and [`Policy::Threshold`] according to odds
     #[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)))
+    fn compile_tr_policy(&self) -> Result<TapTree<Pk>, Error> {
+        let leaf_compilations: Vec<_> = self
+            .flatten_policy(1.0)
+            .into_iter()
+            .map(|(prob, ref policy)| (OrdF64(prob), policy.compile::<Tap>().unwrap()))
+            .collect();
+        let taptree = Self::with_huffman_tree(leaf_compilations, |x| x).unwrap();
+        Ok(taptree)
     }
-    /// Extract the Taproot internal_key from policy tree.
+
+    /// Extract the 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")),
+    fn extract_key(policy: &Self, unspendable_key: Option<Pk>) -> Result<(Pk, Policy<Pk>), Error> {
+        let semantic_policy = policy.lift()?;
+        let concrete_keys = policy.keys().into_iter().collect::<HashSet<_>>();
+        let mut internal_key: Option<Pk> = None;
+
+        for key in concrete_keys.iter() {
+            if semantic_policy
+                .clone()
+                .satisfy_constraint(&Semantic::KeyHash(key.to_pubkeyhash()), true)
+                == Semantic::Trivial
+            {
+                internal_key = Some((*key).clone());
+                break;
+            }
+        }
+
+        match (internal_key, unspendable_key) {
+            (Some(key), _) => Ok((key.clone(), policy.translate_unsatisfiable_pk(&key))),
+            (_, Some(key)) => Ok((key, policy.clone())),
+            _ => Err(errstr("No viable internal key found.")),
         }
     }
 
     /// Compile the [`Tr`] descriptor into optimized [`TapTree`] implementation
+    // TODO: We might require other compile errors for Taproot. Will discuss and update.
     #[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();
+        let tree =
+            Descriptor::new_tr(internal_key, Some(policy.compile_tr_policy().unwrap())).unwrap();
         Ok(tree)
     }
 
@@ -264,6 +349,30 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
+    /// Translate `Semantic::Key(key)` to `Semantic::Unsatisfiable` when extracting TapKey
+    pub fn translate_unsatisfiable_pk(&self, key: &Pk) -> Policy<Pk> {
+        match self {
+            Policy::Key(k) if *k == *key => Policy::Unsatisfiable,
+            Policy::And(ref subs) => Policy::And(
+                subs.iter()
+                    .map(|sub| sub.translate_unsatisfiable_pk(key))
+                    .collect::<Vec<_>>(),
+            ),
+            Policy::Or(ref subs) => Policy::Or(
+                subs.iter()
+                    .map(|(k, sub)| (*k, sub.translate_unsatisfiable_pk(key)))
+                    .collect::<Vec<_>>(),
+            ),
+            Policy::Threshold(k, ref subs) => Policy::Threshold(
+                *k,
+                subs.iter()
+                    .map(|sub| sub.translate_unsatisfiable_pk(key))
+                    .collect::<Vec<_>>(),
+            ),
+            x => x.clone(),
+        }
+    }
+
     /// Get all keys in the policy
     pub fn keys(&self) -> Vec<&Pk> {
         match *self {
diff --git a/src/policy/mod.rs b/src/policy/mod.rs
@@ -374,16 +374,13 @@ mod tests {
     #[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 unspendable_key: String = "UNSPENDABLE".to_string();
+        let policy: Concrete<String> = policy_str!("thresh(2,pk(A),pk(A),pk(C),pk(D))");
+        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);
-        }
+        let ms_compilation: Miniscript<String, Tap> = ms_str!("multi_a(2,C,D)");
+        let tree: TapTree<String> = TapTree::Leaf(Arc::new(ms_compilation));
+        let expected_descriptor = Descriptor::new_tr("A".to_string(), Some(tree)).unwrap();
+        assert_eq!(descriptor, expected_descriptor);
     }
 }
diff --git a/src/policy/semantic.rs b/src/policy/semantic.rs
@@ -189,7 +189,7 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
     // policy.
     // Witness is currently encoded as policy. Only accepts leaf fragment and
     // a normalized policy
-    fn satisfy_constraint(self, witness: &Policy<Pk>, available: bool) -> Policy<Pk> {
+    pub(crate) fn satisfy_constraint(self, witness: &Policy<Pk>, available: bool) -> Policy<Pk> {
         debug_assert!(self.clone().normalized() == self.clone());
         match *witness {
             // only for internal purposes, safe to use unreachable!
@@ -679,10 +679,7 @@ mod tests {
         let policy = StringPolicy::from_str("or(pkh(),older(1000))").unwrap();
         assert_eq!(
             policy,
-            Policy::Threshold(
-                1,
-                vec![Policy::KeyHash("".to_owned()), Policy::Older(1000),]
-            )
+            Policy::Threshold(1, vec![Policy::KeyHash("".to_owned()), Policy::Older(1000)])
         );
         assert_eq!(policy.relative_timelocks(), vec![1000]);
         assert_eq!(policy.absolute_timelocks(), vec![]);
@@ -698,7 +695,7 @@ mod tests {
             policy,
             Policy::Threshold(
                 1,
-                vec![Policy::KeyHash("".to_owned()), Policy::Unsatisfiable,]
+                vec![Policy::KeyHash("".to_owned()), Policy::Unsatisfiable],
             )
         );
         assert_eq!(policy.relative_timelocks(), vec![]);
@@ -711,7 +708,7 @@ mod tests {
             policy,
             Policy::Threshold(
                 2,
-                vec![Policy::KeyHash("".to_owned()), Policy::Unsatisfiable,]
+                vec![Policy::KeyHash("".to_owned()), Policy::Unsatisfiable],
             )
         );
         assert_eq!(policy.relative_timelocks(), vec![]);
@@ -735,7 +732,7 @@ mod tests {
                     Policy::Older(1000),
                     Policy::Older(2000),
                     Policy::Older(2000),
-                ]
+                ],
             )
         );
         assert_eq!(
@@ -759,7 +756,7 @@ mod tests {
                     Policy::Older(1000),
                     Policy::Unsatisfiable,
                     Policy::Unsatisfiable,
-                ]
+                ],
             )
         );
         assert_eq!(
