Add taproot compiler default version

The `compile_tr` method added here uses the heuristic as specified in the docs and provides better cost  guarantees than the `compile_tr_private` method.

Author: SarcasticNastik

src/policy/compiler.rs

@@ -1128,6 +1128,20 @@ pub fn best_compilation<Pk: MiniscriptKey, Ctx: ScriptContext>(
     }
 }
 
+/// Obtain the best compilation of for p=1.0 and q=0, along with the satisfaction cost for the script
+pub(crate) fn best_compilation_sat<Pk: MiniscriptKey, Ctx: ScriptContext>(
+    policy: &Concrete<Pk>,
+) -> Result<(Arc<Miniscript<Pk, Ctx>>, f64), CompilerError> {
+    let mut policy_cache = PolicyCache::<Pk, Ctx>::new();
+    let x: AstElemExt<Pk, Ctx> = best_t(&mut policy_cache, policy, 1.0, None)?;
+    if !x.ms.ty.mall.safe {
+        Err(CompilerError::TopLevelNonSafe)
+    } else if !x.ms.ty.mall.non_malleable {
+        Err(CompilerError::ImpossibleNonMalleableCompilation)
+    } else {
+        Ok((x.ms, x.comp_ext_data.sat_cost))
+    }
+}
 /// Obtain the best B expression with given sat and dissat
 fn best_t<Pk, Ctx>(
     policy_cache: &mut PolicyCache<Pk, Ctx>,

src/policy/concrete.rs

@@ -31,6 +31,7 @@ use {
     crate::Miniscript,
     crate::Tap,
     std::cmp::Reverse,
+    std::collections::BTreeMap,
     std::collections::{BinaryHeap, HashMap},
     std::sync::Arc,
 };
@@ -41,6 +42,14 @@ use crate::miniscript::limits::{LOCKTIME_THRESHOLD, SEQUENCE_LOCKTIME_TYPE_FLAG}
 use crate::miniscript::types::extra_props::TimelockInfo;
 use crate::{errstr, Error, ForEach, ForEachKey, MiniscriptKey};
 
+/// [`TapTree`] -> ([`Policy`], satisfaction cost) cache
+#[cfg(feature = "compiler")]
+type PolicyTapCache<Pk> = BTreeMap<TapTree<Pk>, (Policy<Pk>, f64)>;
+
+/// [`Miniscript`] -> leaf probability in policy cache
+#[cfg(feature = "compiler")]
+type MsTapCache<Pk> = BTreeMap<TapTree<Pk>, f64>;
+
 /// Concrete policy which corresponds directly to a Miniscript structure,
 /// and whose disjunctions are annotated with satisfaction probabilities
 /// to assist the compiler
@@ -263,6 +272,62 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
+    /// Compile [`Policy`] into a [`TapTree Descriptor`][`Descriptor::Tr`]
+    ///
+    ///
+    ///  This follows the heuristic as described in [`with_huffman_tree_eff`]
+    #[cfg(feature = "compiler")]
+    pub fn compile_tr(&self, unspendable_key: Option<Pk>) -> Result<Descriptor<Pk>, Error> {
+        self.is_valid()?; // Check for validity
+        match self.is_safe_nonmalleable() {
+            (false, _) => Err(Error::from(CompilerError::TopLevelNonSafe)),
+            (_, false) => Err(Error::from(
+                CompilerError::ImpossibleNonMalleableCompilation,
+            )),
+            _ => {
+                let (internal_key, policy) = self.clone().extract_key(unspendable_key)?;
+                let tree = Descriptor::new_tr(
+                    internal_key,
+                    match policy {
+                        Policy::Trivial => None,
+                        policy => {
+                            let mut policy_cache = PolicyTapCache::<Pk>::new();
+                            let mut ms_cache = MsTapCache::<Pk>::new();
+                            // Obtain the policy compilations and populate the respective caches for
+                            // creating the huffman tree later on
+                            let leaf_compilations: Vec<_> = policy
+                                .to_tapleaf_prob_vec(1.0)
+                                .into_iter()
+                                .filter(|x| x.1 != Policy::Unsatisfiable)
+                                .map(|(prob, ref pol)| {
+                                    let compilation =
+                                        compiler::best_compilation_sat::<Pk, Tap>(pol).unwrap();
+                                    policy_cache.insert(
+                                        TapTree::Leaf(Arc::clone(&compilation.0)),
+                                        (pol.clone(), compilation.1), // (policy, sat_cost)
+                                    );
+                                    ms_cache.insert(
+                                        TapTree::Leaf(Arc::from(compilation.0.clone())),
+                                        prob,
+                                    );
+                                    compilation.0
+                                })
+                                .collect();
+                            let taptree = with_huffman_tree_eff(
+                                leaf_compilations,
+                                &mut policy_cache,
+                                &mut ms_cache,
+                            )
+                            .unwrap();
+                            Some(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> {
@@ -805,6 +870,36 @@ where
     }
 }
 
+/// Average satisfaction cost for [`TapTree`] with the leaf [`Miniscript`] nodes having
+/// probabilities corresponding to the (sub)policies they're compiled from.
+///
+/// Average satisfaction cost for [`TapTree`] over script-spend paths is probability times
+/// the size of control block + the script size.
+#[cfg(feature = "compiler")]
+fn taptree_cost<Pk: MiniscriptKey>(
+    tr: &TapTree<Pk>,
+    ms_cache: &MsTapCache<Pk>,
+    policy_cache: &PolicyTapCache<Pk>,
+    depth: u32,
+) -> f64 {
+    match *tr {
+        TapTree::Tree(ref l, ref r) => {
+            taptree_cost(l, ms_cache, policy_cache, depth + 1)
+                + taptree_cost(r, ms_cache, policy_cache, depth + 1)
+        }
+        TapTree::Leaf(ref ms) => {
+            let prob = ms_cache
+                .get(&TapTree::Leaf(Arc::clone(ms)))
+                .expect("Probability should exist for the given ms");
+            let sat_cost = policy_cache
+                .get(&TapTree::Leaf(Arc::clone(ms)))
+                .expect("Cost should exist for the given ms")
+                .1;
+            prob * (ms.script_size() as f64 + sat_cost + 32.0 * depth as f64)
+        }
+    }
+}
+
 /// Create a Huffman Tree from compiled [Miniscript] nodes
 #[cfg(feature = "compiler")]
 fn with_huffman_tree<Pk: MiniscriptKey>(
@@ -835,3 +930,117 @@ fn with_huffman_tree<Pk: MiniscriptKey>(
         .1;
     Ok(node)
 }
+
+/// Create a [`TapTree`] from the a list of [`Miniscript`]s having corresponding satisfaction
+/// cost and probability.
+///
+/// Given that satisfaction probability and cost for each script is known, constructing the
+/// [`TapTree`] as a huffman tree over the net cost (as defined in [`taptree_cost`]) is
+/// the optimal one.
+/// For finding the optimal policy to taptree compilation, we are required to search
+/// exhaustively over all policies which have the same leaf policies. Owing to the exponential
+/// blow-up for such a method, we use a heuristic where we augment the merge to check if the
+/// compilation of a new (sub)policy into a [`TapTree::Leaf`] with the policy corresponding to
+/// the nodes as children is better than [`TapTree::Tree`] with the nodes as children.
+#[cfg(feature = "compiler")]
+fn with_huffman_tree_eff<Pk: MiniscriptKey>(
+    ms: Vec<Arc<Miniscript<Pk, Tap>>>,
+    policy_cache: &mut PolicyTapCache<Pk>,
+    ms_cache: &mut MsTapCache<Pk>,
+) -> Result<TapTree<Pk>, Error> {
+    let mut node_weights = BinaryHeap::<(Reverse<OrdF64>, OrdF64, TapTree<Pk>)>::new(); // (cost, branch_prob, tree)
+                                                                                        // Populate the heap with each `ms` as a TapLeaf, and the respective cost fields
+    for script in ms {
+        let wt = OrdF64(taptree_cost(
+            &TapTree::Leaf(Arc::clone(&script)),
+            ms_cache,
+            policy_cache,
+            0,
+        ));
+        let prob = OrdF64(
+            *ms_cache
+                .get(&TapTree::Leaf(Arc::clone(&script)))
+                .expect("Probability should exist for the given ms"),
+        );
+        node_weights.push((Reverse(wt), prob, TapTree::Leaf(Arc::clone(&script))));
+    }
+    if node_weights.is_empty() {
+        return Err(errstr("Empty Miniscript compilation"));
+    }
+    while node_weights.len() > 1 {
+        // Obtain the two least-weighted nodes from the heap for merging
+        let (_prev_cost1, p1, ms1) = node_weights.pop().expect("len must atleast be two");
+        let (_prev_cost2, p2, ms2) = node_weights.pop().expect("len must atleast be two");
+
+        // Retrieve the respective policies
+        let (left_pol, _c1) = policy_cache
+            .get(&ms1)
+            .ok_or_else(|| errstr("No corresponding policy found"))?
+            .clone();
+
+        let (right_pol, _c2) = policy_cache
+            .get(&ms2)
+            .ok_or_else(|| errstr("No corresponding policy found"))?
+            .clone();
+
+        // Create a parent policy with the respective node TapTrees as children (with odds
+        // weighted approximately in ratio to their probabilities)
+        let parent_policy = Policy::Or(vec![
+            ((p1.0 * 1e4).round() as usize, left_pol),
+            ((p2.0 * 1e4).round() as usize, right_pol),
+        ]);
+
+        // Obtain compilation for the parent policy
+        let (parent_compilation, parent_sat_cost) =
+            compiler::best_compilation_sat::<Pk, Tap>(&parent_policy)?;
+
+        // Probability of the parent node being satisfied equals the probability of either
+        // nodes to be satisfied. Since we weight the odds appropriately, the children nodes
+        // still have approximately the same probabilities
+        let p = p1.0 + p2.0;
+        // Inserting parent policy's weights (sat_cost and probability) for later usage
+        ms_cache.insert(TapTree::Leaf(Arc::clone(&parent_compilation)), p);
+        policy_cache.insert(
+            TapTree::Leaf(Arc::clone(&parent_compilation)),
+            (parent_policy.clone(), parent_sat_cost),
+        );
+
+        let parent_cost = OrdF64(taptree_cost(
+            &TapTree::Leaf(Arc::clone(&parent_compilation)),
+            ms_cache,
+            policy_cache,
+            0,
+        ));
+        let children_cost = OrdF64(
+            taptree_cost(&ms1, ms_cache, policy_cache, 0)
+                + taptree_cost(&ms2, ms_cache, policy_cache, 0),
+        );
+
+        // Merge the children nodes into either TapLeaf of the parent compilation or
+        // TapTree children nodes accordingly
+        node_weights.push(if parent_cost > children_cost {
+            ms_cache.insert(
+                TapTree::Tree(Arc::from(ms1.clone()), Arc::from(ms2.clone())),
+                p,
+            );
+            policy_cache.insert(
+                TapTree::Tree(Arc::from(ms1.clone()), Arc::from(ms2.clone())),
+                (parent_policy, parent_sat_cost),
+            );
+            (
+                Reverse(children_cost),
+                OrdF64(p),
+                TapTree::Tree(Arc::from(ms1), Arc::from(ms2)),
+            )
+        } else {
+            let node = TapTree::Leaf(Arc::from(parent_compilation));
+            (Reverse(parent_cost), OrdF64(p), node)
+        });
+    }
+    debug_assert!(node_weights.len() == 1);
+    let node = node_weights
+        .pop()
+        .expect("huffman tree algorithm is broken")
+        .2;
+    Ok(node)
+}