trivial P2Tr compiler done

Taproot internal-key extraction from policy done

Author: SarcasticNastik

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(crate) struct OrdF64(pub f64);
 
 impl Eq for OrdF64 {}
 impl Ord for OrdF64 {

src/policy/concrete.rs

@@ -27,8 +27,15 @@ use miniscript::limits::{HEIGHT_TIME_THRESHOLD, SEQUENCE_LOCKTIME_TYPE_FLAG};
 use miniscript::types::extra_props::TimeLockInfo;
 #[cfg(feature = "compiler")]
 use {
-    descriptor::TapTree, miniscript::ScriptContext, policy::compiler,
-    policy::compiler::CompilerError, std::sync::Arc, Descriptor, Miniscript, Tap,
+    descriptor::TapTree,
+    miniscript::ScriptContext,
+    policy::compiler::{CompilerError, OrdF64},
+    policy::Concrete,
+    policy::{compiler, Liftable, Semantic},
+    std::cmp::Reverse,
+    std::collections::{BinaryHeap, HashMap},
+    std::sync::Arc,
+    Descriptor, Miniscript, Tap,
 };
 use {Error, ForEach, ForEachKey, MiniscriptKey};
 
@@ -126,27 +133,137 @@ 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 compile_leaf_taptree(&self) -> Result<TapTree<Pk>, Error> {
-        let compilation = self.compile::<Tap>().unwrap();
-        Ok(TapTree::Leaf(Arc::new(compilation)))
+    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>, TapTree<Pk>)>::new();
+        for (prob, script) in ms {
+            node_weights.push((Reverse(f(prob)), 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)),
+                TapTree::Tree(Arc::from(s1), Arc::from(s2)),
+            ));
+        }
+
+        debug_assert!(node_weights.len() == 1);
+        let node = node_weights
+            .pop()
+            .expect("huffman tree algorithm is broken")
+            .1;
+        Ok(node)
     }
 
-    /// Extract the Taproot internal_key from policy tree.
+    /// 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 extract_key(&self, unspendable_key: Option<Pk>) -> Result<(Pk, &Policy<Pk>), Error> {
-        match unspendable_key {
-            Some(key) => Ok((key, self)),
-            None => Err(errstr("No internal key found")),
+    fn to_tapleaf_prob_vec(&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.to_tapleaf_prob_vec(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.to_tapleaf_prob_vec(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_tr_policy(&self) -> Result<TapTree<Pk>, Error> {
+        let leaf_compilations: Vec<_> = self
+            .to_tapleaf_prob_vec(1.0)
+            .into_iter()
+            .filter(|x| x.1 != Policy::Unsatisfiable)
+            .map(|(prob, ref policy)| (OrdF64(prob), compiler::best_compilation(policy).unwrap()))
+            .collect();
+        let taptree = Self::with_huffman_tree(leaf_compilations, |x| x).unwrap();
+        Ok(taptree)
+    }
+
+    /// Extract the internal_key from policy tree.
+    #[cfg(feature = "compiler")]
+    fn extract_key(self, unspendable_key: Option<Pk>) -> Result<(Pk, Policy<Pk>), Error> {
+        // Making sure the borrow ends before you move the value.
+        let mut internal_key: Option<Pk> = None;
+        {
+            let mut prob = 0.;
+            let semantic_policy = self.lift()?;
+            let concrete_keys = self.keys();
+            let key_prob_map: HashMap<_, _> = self
+                .to_tapleaf_prob_vec(1.0)
+                .into_iter()
+                .filter(|(_, ref pol)| match *pol {
+                    Concrete::Key(..) => true,
+                    _ => false,
+                })
+                .map(|(prob, key)| (key, prob))
+                .collect();
+
+            for key in concrete_keys.into_iter() {
+                if semantic_policy
+                    .clone()
+                    .satisfy_constraint(&Semantic::KeyHash(key.to_pubkeyhash()), true)
+                    == Semantic::Trivial
+                {
+                    match key_prob_map.get(&Concrete::Key(key.clone())) {
+                        Some(val) => {
+                            if *val > prob {
+                                prob = *val;
+                                internal_key = Some(key.clone());
+                            }
+                        }
+                        None => return Err(errstr("Key should have existed in the HashMap!")),
+                    }
+                }
+            }
+        }
+        match (internal_key, unspendable_key) {
+            (Some(ref key), _) => Ok((key.clone(), self.translate_unsatisfiable_pk(&key))),
+            (_, Some(key)) => Ok((key, self)),
+            _ => 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(unspendable_key)?;
-        let tree = Descriptor::new_tr(internal_key, Some(policy.compile_leaf_taptree()?))?;
+        let (internal_key, policy) = self.clone().extract_key(unspendable_key)?;
+        let tree = Descriptor::new_tr(
+            internal_key,
+            match policy {
+                Policy::Trivial => None,
+                policy => Some(policy.compile_tr_policy()?),
+            },
+        )?;
         Ok(tree)
     }
 
@@ -250,6 +367,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(ref k) if k.clone() == *key => Policy::Unsatisfiable,
+            Policy::And(subs) => Policy::And(
+                subs.into_iter()
+                    .map(|sub| sub.translate_unsatisfiable_pk(key))
+                    .collect::<Vec<_>>(),
+            ),
+            Policy::Or(subs) => Policy::Or(
+                subs.into_iter()
+                    .map(|(k, sub)| (k, sub.translate_unsatisfiable_pk(key)))
+                    .collect::<Vec<_>>(),
+            ),
+            Policy::Threshold(k, subs) => Policy::Threshold(
+                k,
+                subs.into_iter()
+                    .map(|sub| sub.translate_unsatisfiable_pk(key))
+                    .collect::<Vec<_>>(),
+            ),
+            x => x,
+        }
+    }
+
     /// Get all keys in the policy
     pub fn keys(&self) -> Vec<&Pk> {
         match *self {

src/policy/mod.rs

@@ -21,26 +21,26 @@
 //! The format represents EC public keys abstractly to allow wallets to replace
 //! these with BIP32 paths, pay-to-contract instructions, etc.
 //!
-use {error, fmt};
-
-#[cfg(feature = "compiler")]
-pub mod compiler;
-pub mod concrete;
-pub mod semantic;
-
 use descriptor::Descriptor;
 use miniscript::{Miniscript, ScriptContext};
+use Error;
+use MiniscriptKey;
 use Terminal;
+use {error, fmt};
 
 pub use self::concrete::Policy as Concrete;
 /// Semantic policies are "abstract" policies elsewhere; but we
 /// avoid this word because it is a reserved keyword in Rust
 pub use self::semantic::Policy as Semantic;
-use Error;
-use MiniscriptKey;
+
+#[cfg(feature = "compiler")]
+pub mod compiler;
+pub mod concrete;
+pub mod semantic;
 
 /// Policy entailment algorithm maximum number of terminals allowed
 const ENTAILMENT_MAX_TERMINALS: usize = 20;
+
 /// Trait describing script representations which can be lifted into
 /// an abstract policy, by discarding information.
 /// After Lifting all policies are converted into `KeyHash(Pk::HasH)` to
@@ -228,20 +228,23 @@ impl<Pk: MiniscriptKey> Liftable<Pk> for Concrete<Pk> {
 
 #[cfg(test)]
 mod tests {
-    use super::{
-        super::miniscript::{context::Segwitv0, Miniscript},
-        Concrete, Liftable, Semantic,
-    };
-    use bitcoin;
-    #[cfg(feature = "compiler")]
-    use descriptor::TapTree;
     use std::str::FromStr;
     #[cfg(feature = "compiler")]
     use std::sync::Arc;
+
+    use bitcoin;
+
+    #[cfg(feature = "compiler")]
+    use descriptor::TapTree;
     use DummyKey;
     #[cfg(feature = "compiler")]
     use {Descriptor, Tap};
 
+    use super::{
+        super::miniscript::{context::Segwitv0, Miniscript},
+        Concrete, Liftable, Semantic,
+    };
+
     type ConcretePol = Concrete<DummyKey>;
     type SemanticPol = Semantic<DummyKey>;
 
@@ -272,6 +275,7 @@ mod tests {
         // thresh with k = 2
         assert!(ConcretePol::from_str("thresh(2,after(1000000000),after(100),pk())").is_err());
     }
+
     #[test]
     fn policy_rtt_tests() {
         concrete_policy_rtt("pk()");
@@ -361,28 +365,46 @@ mod tests {
                         2,
                         vec![
                             Semantic::KeyHash(key_a.pubkey_hash().as_hash()),
-                            Semantic::Older(42)
-                        ]
+                            Semantic::Older(42),
+                        ],
                     ),
-                    Semantic::KeyHash(key_b.pubkey_hash().as_hash())
-                ]
+                    Semantic::KeyHash(key_b.pubkey_hash().as_hash()),
+                ],
             ),
             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);
+    fn taproot_compile() {
+        // Trivial single-node compilation
+        let unspendable_key: String = "UNSPENDABLE".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({},A,B,C,D)", k);
+            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();
+            let expected_descriptor =
+                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
+            assert_eq!(descriptor, expected_descriptor);
+        }
+
+        // Trivial multi-node compilation
+        {
+            let policy: Concrete<String> = policy_str!("or(and(pk(A),pk(B)),and(pk(C),pk(D)))");
+            let descriptor = policy.compile_tr(Some(unspendable_key.clone())).unwrap();
 
+            let left_ms_compilation: Arc<Miniscript<String, Tap>> =
+                Arc::new(ms_str!("and_v(v:pk(C),pk(D))"));
+            let right_ms_compilation: Arc<Miniscript<String, Tap>> =
+                Arc::new(ms_str!("and_v(v:pk(A),pk(B))"));
+            let left_node: Arc<TapTree<String>> = Arc::from(TapTree::Leaf(left_ms_compilation));
+            let right_node: Arc<TapTree<String>> = Arc::from(TapTree::Leaf(right_ms_compilation));
+            let tree: TapTree<String> = TapTree::Tree(left_node, right_node);
+            let expected_descriptor =
+                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
             assert_eq!(descriptor, expected_descriptor);
         }
     }

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!