Add P2Tr compiler

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

Author: SarcasticNastik

src/policy/compiler.rs

@@ -36,7 +36,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

@@ -25,11 +25,13 @@ use {
     crate::descriptor::TapTree,
     crate::miniscript::ScriptContext,
     crate::policy::compiler::CompilerError,
+    crate::policy::compiler::OrdF64,
     crate::policy::{compiler, Concrete, Liftable, Semantic},
     crate::Descriptor,
     crate::Miniscript,
     crate::Tap,
-    std::collections::HashMap,
+    std::cmp::Reverse,
+    std::collections::{BinaryHeap, HashMap},
     std::sync::Arc,
 };
 
@@ -173,15 +175,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.;
@@ -222,11 +232,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)
     }
 
@@ -771,3 +798,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

@@ -368,15 +368,35 @@ mod tests {
 
     #[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, Some(tree)).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);
+        }
+
+        // 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);
+        }
     }
 }