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

@@ -1117,6 +1117,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

@@ -33,12 +33,21 @@ use {
     policy::Concrete,
     policy::{compiler, Liftable, Semantic},
     std::cmp::Reverse,
+    std::collections::BTreeMap,
     std::collections::{BinaryHeap, HashMap},
     std::sync::Arc,
     Descriptor, Miniscript, Tap,
 };
 use {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
@@ -169,9 +178,154 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         Ok(node)
     }
 
-    /// 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)
+    /// 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(
+        tr: &TapTree<Pk>,
+        ms_cache: &MsTapCache<Pk>,
+        policy_cache: &PolicyTapCache<Pk>,
+        depth: u32,
+    ) -> f64 {
+        match *tr {
+            TapTree::Tree(ref l, ref r) => {
+                Self::taptree_cost(l, ms_cache, policy_cache, depth + 1)
+                    + Self::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 [`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 [`Policy::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(
+        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(Self::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(Self::taptree_cost(
+                &TapTree::Leaf(Arc::clone(&parent_compilation)),
+                ms_cache,
+                policy_cache,
+                0,
+            ));
+            let children_cost = OrdF64(
+                Self::taptree_cost(&ms1, ms_cache, policy_cache, 0)
+                    + Self::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)
+    }
+
+    /// Flatten the [`Policy`] tree structure into a Vector of (sub)-policies with corresponding
+    /// leaf probabilities by using odds for calculating branch probabilities.
+    /// We consider splitting the [`Policy`] tree into respective (sub)-policies considering
+    /// disjunction over [`Policy::Or`] and [`Policy::Threshold`](1, ...)
     #[cfg(feature = "compiler")]
     fn to_tapleaf_prob_vec(&self, prob: f64) -> Vec<(f64, Policy<Pk>)> {
         match *self {
@@ -280,6 +434,62 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
+    /// Compile [`Policy`] into a [`TapTree Descriptor`][`Descriptor::Tr`]
+    ///
+    ///
+    ///  This follows the heuristic as described in [`Policy::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 = Self::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> {