Add P2Tr compiler

Introduce a `compiler_tr` API for compiling a policy to a Tr Descriptor.

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

@@ -29,10 +29,11 @@ use miniscript::types::extra_props::TimeLockInfo;
 use {
     descriptor::TapTree,
     miniscript::ScriptContext,
-    policy::compiler::CompilerError,
+    policy::compiler::{CompilerError, OrdF64},
     policy::Concrete,
     policy::{compiler, Liftable, Semantic},
-    std::collections::HashMap,
+    std::cmp::Reverse,
+    std::collections::{BinaryHeap, HashMap},
     std::sync::Arc,
     Descriptor, Miniscript, Tap,
 };
@@ -156,15 +157,23 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
+    /// Compile [`Policy::Or`] and [`Policy::Threshold`] according to odds
     #[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 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 = with_huffman_tree::<Pk>(leaf_compilations).unwrap();
+        Ok(taptree)
     }
 
-    /// Extract the Taproot internal_key from policy tree.
+    /// 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.;
@@ -205,11 +214,28 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
-    /// Compile the [`Tr`] descriptor into optimized [`TapTree`] implementation
+    /// Compile the [`Policy`] into a [`Tr`][`Descriptor::Tr`] Descriptor.
+    ///
+    /// ### TapTree compilation
+    ///
+    /// The policy tree constructed by root-level disjunctions over [`Or`][`Policy::Or`] and
+    /// [`Thresh`][`Policy::Threshold`](1, ..) which is flattened into a vector (with respective
+    /// probabilities derived from odds) of policies.
+    /// For example, the policy `thresh(1,or(pk(A),pk(B)),and(or(pk(C),pk(D)),pk(E)))` gives the vector
+    /// `[pk(A),pk(B),and(or(pk(C),pk(D)),pk(E)))]`. Each policy in the vector is compiled into
+    /// the respective miniscripts. A Huffman Tree is created from this vector which optimizes over
+    /// the probabilitity of satisfaction for the respective branch in the TapTree.
+    // TODO: We might require other compile errors for Taproot.
     #[cfg(feature = "compiler")]
     pub fn compile_tr(&self, unspendable_key: Option<Pk>) -> Result<Descriptor<Pk>, Error> {
         let (internal_key, policy) = self.clone().extract_key(unspendable_key)?;
-        let tree = Descriptor::new_tr(internal_key, Some(policy.compile_leaf_taptree()?))?;
+        let tree = Descriptor::new_tr(
+            internal_key,
+            match policy {
+                Policy::Trivial => None,
+                policy => Some(policy.compile_tr_policy()?),
+            },
+        )?;
         Ok(tree)
     }
 
@@ -763,3 +789,34 @@ where
         Policy::from_tree_prob(top, false).map(|(_, result)| result)
     }
 }
+
+/// Create a Huffman Tree from compiled [Miniscript] nodes
+#[cfg(feature = "compiler")]
+fn with_huffman_tree<Pk: MiniscriptKey>(
+    ms: Vec<(OrdF64, Miniscript<Pk, Tap>)>,
+) -> Result<TapTree<Pk>, Error> {
+    let mut node_weights = BinaryHeap::<(Reverse<OrdF64>, TapTree<Pk>)>::new();
+    for (prob, script) in ms {
+        node_weights.push((Reverse(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)
+}

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>;
 
@@ -361,27 +364,47 @@ 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() {
-        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();
+    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(2,A,B,C,D)");
+            let tree: TapTree<String> = TapTree::Leaf(Arc::new(ms_compilation));
+            let expected_descriptor =
+                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
+            assert_eq!(descriptor, expected_descriptor);
+        }
 
-        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();
+        // 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();
 
-        assert_eq!(descriptor, expected_descriptor);
+            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);
+        }
     }
 }