Implement tr-compiler-v3 enumeration compiler

Author: SarcasticNastik

src/lib.rs

@@ -985,7 +985,7 @@ mod prelude {
     pub use alloc::{
         borrow::{Borrow, Cow, ToOwned},
         boxed::Box,
-        collections::{vec_deque::VecDeque, BTreeMap, BinaryHeap},
+        collections::{vec_deque::VecDeque, BTreeMap, BTreeSet, BinaryHeap},
         rc, slice,
         string::{String, ToString},
         sync,
@@ -995,7 +995,7 @@ mod prelude {
     pub use std::{
         borrow::{Borrow, Cow, ToOwned},
         boxed::Box,
-        collections::{vec_deque::VecDeque, BTreeMap, BinaryHeap, HashMap, HashSet},
+        collections::{vec_deque::VecDeque, BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet},
         rc, slice,
         string::{String, ToString},
         sync,

src/miniscript/limits.rs

@@ -50,3 +50,6 @@ pub const MAX_BLOCK_WEIGHT: usize = 4000000;
 /// Maximum pubkeys as arguments to CHECKMULTISIG
 // https://github.com/bitcoin/bitcoin/blob/6acda4b00b3fc1bfac02f5de590e1a5386cbc779/src/script/script.h#L30
 pub const MAX_PUBKEYS_PER_MULTISIG: usize = 20;
+
+/// Maximum TapLeafs allowed in a compiled TapTree
+pub const MAX_TAPLEAFS_PER_TAPTREE: usize = 10_000;

src/policy/concrete.rs

@@ -24,6 +24,7 @@ use bitcoin::hashes::{hash160, ripemd160};
 #[cfg(feature = "compiler")]
 use {
     crate::descriptor::TapTree,
+    crate::miniscript::limits::MAX_TAPLEAFS_PER_TAPTREE,
     crate::miniscript::ScriptContext,
     crate::policy::compiler::CompilerError,
     crate::policy::compiler::OrdF64,
@@ -186,9 +187,14 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
     ///                                    B         C
     ///
     /// gives the vector [(2/3, A), (1/3 * 3/4, B), (1/3 * 1/4, C)].
+    ///
+    /// ## Constraints
+    ///
+    /// Since this splitting might lead to exponential blow-up, we constraint the number of
+    /// leaf-nodes to [`MAX_TAPLEAFS_PER_TAPTREE`].
     #[cfg(feature = "compiler")]
     fn to_tapleaf_prob_vec(&self, prob: f64) -> Vec<(f64, Policy<Pk>)> {
-        match *self {
+        match self {
             Policy::Or(ref subs) => {
                 let total_odds: usize = subs.iter().map(|(ref k, _)| k).sum();
                 subs.iter()
@@ -198,17 +204,141 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
                     .flatten()
                     .collect::<Vec<_>>()
             }
-            Policy::Threshold(k, ref subs) if k == 1 => {
+            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())],
+            // Instead of the fixed-point algorithm, consider direct k-choose-n split if limits
+            // don't exceed
+            x => vec![(prob, x.clone())],
+        }
+    }
+
+    /// Given a [`Policy`], return a vector of policies whose disjunction is isomorphic to the initial one.
+    /// This function is supposed to incrementally expand i.e. represent the policy as disjunction over
+    /// sub-policies output by it. The probability calculations are similar as
+    /// [to_tapleaf_prob_vec][`Policy::to_tapleaf_prob_vec`]
+    #[cfg(feature = "compiler")]
+    fn enumerate_pol(&self, prob: f64) -> Vec<(f64, Policy<Pk>)> {
+        match self {
+            Policy::Or(subs) => {
+                let total_odds = subs.iter().fold(0, |acc, x| acc + x.0);
+                subs.iter()
+                    .map(|(odds, pol)| (prob * *odds as f64 / total_odds as f64, pol.clone()))
+                    .collect::<Vec<_>>()
+            }
+            Policy::Threshold(k, subs) if *k == 1 => {
+                let total_odds = subs.len();
+                subs.iter()
+                    .map(|pol| (prob / total_odds as f64, pol.clone()))
+                    .collect::<Vec<_>>()
+            }
+            Policy::Threshold(k, subs) if *k != subs.len() => generate_combination(subs, prob, *k),
+            pol => vec![(prob, pol.clone())],
         }
     }
 
+    /// Generates a root-level disjunctive tree over the given policy tree, by using fixed-point
+    /// algorithm to enumerate the disjunctions until exhaustive root-level enumeration or limits
+    /// exceed.
+    /// For a given [policy][`Policy`], we maintain an [ordered set][`BTreeSet`] of `(prob, policy)`
+    /// (ordered by probability) to maintain the list of enumerated sub-policies whose disjunction
+    /// is isomorphic to initial policy (*invariant*).
+    #[cfg(feature = "compiler")]
+    fn enumerate_policy_tree(&self, prob: f64) -> Vec<(f64, Self)> {
+        let mut tapleaf_prob_vec = BTreeSet::<(Reverse<OrdF64>, Self)>::new();
+        // Store probability corresponding to policy in the enumerated tree. This is required since
+        // owing to the current [policy element enumeration algorithm][`Policy::enumerate_pol`],
+        // two passes of the algorithm might result in same sub-policy showing up. Currently, we
+        // merge the nodes by adding up the corresponding probabilities for the same policy.
+        let mut pol_prob_map = HashMap::<Self, OrdF64>::new();
+
+        tapleaf_prob_vec.insert((Reverse(OrdF64(prob)), self.clone()));
+        pol_prob_map.insert(self.clone(), OrdF64(prob));
+
+        // Since we know that policy enumeration *must* result in increase in total number of nodes,
+        // we can maintain the length of the ordered set to check if the
+        // [enumeration pass][`Policy::enumerate_pol`] results in further policy split or not.
+        let mut prev_len = 0usize;
+        // This is required since we merge some corresponding policy nodes, so we can explicitly
+        // store the variables
+        let mut enum_len = tapleaf_prob_vec.len();
+
+        let mut ret: Vec<(f64, Self)> = vec![];
+
+        // Stopping condition: When NONE of the inputs can be further enumerated.
+        'outer: loop {
+            //--- FIND a plausible node ---
+
+            let mut prob: Reverse<OrdF64> = Reverse(OrdF64(0.0));
+            let mut curr_policy: &Self = &Policy::Unsatisfiable;
+            let mut curr_pol_replace_vec: Vec<(f64, Self)> = vec![(prob.0 .0, curr_policy.clone())];
+
+            // The nodes which can't be enumerated further are directly appended to ret and removed
+            // from the ordered set.
+            let mut to_del: Vec<(f64, Self)> = vec![];
+            for (i, (p, pol)) in tapleaf_prob_vec.iter().enumerate() {
+                curr_pol_replace_vec = pol.enumerate_pol(p.0 .0);
+                enum_len += curr_pol_replace_vec.len() - 1; // A disjunctive node should have seperated this into more nodes
+                if prev_len < enum_len {
+                    // Plausible node found
+                    prob = *p;
+                    curr_policy = pol;
+                    break;
+                } else if i == tapleaf_prob_vec.len() - 1 {
+                    // No enumerable node found i.e. STOP
+                    // Move all the elements to final return set
+                    ret.append(&mut to_del);
+                    ret.push((p.0 .0, pol.clone()));
+                    break 'outer;
+                } else {
+                    // Either node is enumerable, or we have
+                    // Mark all non-enumerable nodes to remove
+                    to_del.push((p.0 .0, pol.clone()));
+                }
+            }
+
+            // --- Sanity Checks ---
+            if enum_len > MAX_TAPLEAFS_PER_TAPTREE || curr_policy == &Policy::Unsatisfiable {
+                break;
+            }
+
+            // If total number of nodes are in limits, we remove the current node and replace it
+            // with children nodes
+
+            // Remove current node
+            tapleaf_prob_vec.remove(&(prob, curr_policy.clone()));
+            // Remove marked nodes (minor optimization)
+            for (p, pol) in to_del {
+                tapleaf_prob_vec.remove(&(Reverse(OrdF64(p)), pol.clone()));
+                ret.push((p, pol.clone()));
+            }
+
+            // Append node if not previously exists, else update the respective probability
+            for (p, policy) in curr_pol_replace_vec {
+
+                match pol_prob_map.get(&policy) {
+                    Some(prev_prob) => {
+                        tapleaf_prob_vec.remove(&(Reverse(*prev_prob), policy.clone()));
+                        tapleaf_prob_vec.insert((Reverse(OrdF64(prev_prob.0 + p)), policy.clone()));
+                        pol_prob_map.insert(policy.clone(), OrdF64(prev_prob.0 + p));
+                    }
+                    None => {
+                        tapleaf_prob_vec.insert((Reverse(OrdF64(p)), policy.clone()));
+                        pol_prob_map.insert(policy.clone(), OrdF64(p));
+                    }
+                }
+            }
+            // --- Update --- total sub-policies count (considering no merging of nodes)
+            prev_len = enum_len;
+        }
+
+        ret
+    }
+
     /// Extract the internal_key from policy tree.
     #[cfg(feature = "compiler")]
     fn extract_key(self, unspendable_key: Option<Pk>) -> Result<(Pk, Policy<Pk>), Error> {
@@ -305,6 +435,56 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
+    /// Compile the [`Policy`] into a [`Tr`][`Descriptor::Tr`] Descriptor, with policy-enumeration
+    /// by [`Policy::enumerate_policy_tree`].
+    ///
+    /// ### TapTree compilation
+    ///
+    /// The policy tree constructed by root-level disjunctions over [`Or`][`Policy::Or`] and
+    /// [`Thresh`][`Policy::Threshold`](k, ..n..) 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)))]`.
+    ///
+    /// ### Policy enumeration
+    ///
+    /// Refer to [`Policy::enumerate_policy_tree`] for the current strategy implemented.
+    #[cfg(feature = "compiler")]
+    pub fn compile_tr_private_experimental(
+        &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 leaf_compilations: Vec<_> = policy
+                                .enumerate_policy_tree(1.0)
+                                .into_iter()
+                                .filter(|x| x.1 != Policy::Unsatisfiable)
+                                .map(|(prob, ref pol)| {
+                                    (OrdF64(prob), compiler::best_compilation(pol).unwrap())
+                                })
+                                .collect();
+                            let taptree = with_huffman_tree::<Pk>(leaf_compilations).unwrap();
+                            Some(taptree)
+                        }
+                    },
+                )?;
+                Ok(tree)
+            }
+        }
+    }
+
     /// Compile the [`Policy`] into desc_ctx [`Descriptor`]
     ///
     /// In case of [Tr][`DescriptorCtx::Tr`], `internal_key` is used for the Taproot comilation when
@@ -508,15 +688,17 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
     fn num_tap_leaves(&self) -> usize {
         match self {
             Policy::Or(subs) => subs.iter().map(|(_prob, pol)| pol.num_tap_leaves()).sum(),
-            Policy::Threshold(k, subs) if *k == 1 => combination(subs.len(), *k),
+            Policy::Threshold(k, subs) if *k == 1 => {
+                subs.iter().map(|pol| pol.num_tap_leaves()).sum()
+            }
             _ => 1,
         }
     }
 
     /// Check on the number of TapLeaves
     #[cfg(feature = "compiler")]
-     fn check_tapleaf(&self) -> Result<(), Error> {
-        if self.num_tap_leaves() > 100_000 {
+    fn check_tapleaf(&self) -> Result<(), Error> {
+        if self.num_tap_leaves() > MAX_TAPLEAFS_PER_TAPTREE {
             return Err(errstr("Too many Tapleaves"));
         }
         Ok(())
@@ -939,12 +1121,32 @@ fn with_huffman_tree<Pk: MiniscriptKey>(
         .pop()
         .expect("huffman tree algorithm is broken")
         .1;
+    Ok(node)
+}
 
+/// Enumerate a [Thresh][`Policy::Threshold`](k, ..n..) into `n` different thresh.
+///
+/// ## Strategy
+///
+/// `thresh(k, x_1...x_n) := thresh(1, thresh(k, x_2...x_n), thresh(k, x_1x_3...x_n), ...., thresh(k, x_1...x_{n-1}))`
+/// by the simple argument that choosing `k` conditions from `n` available conditions might not contain
+/// any one of the conditions exclusively.
 #[cfg(feature = "compiler")]
-fn combination(n: usize, k: usize) -> usize {
-    if k == 0 {
-        1
-    } else {
-        n * combination(n - 1, k - 1) / k
+fn generate_combination<Pk: MiniscriptKey>(
+    policy_vec: &Vec<Policy<Pk>>,
+    prob: f64,
+    k: usize,
+) -> Vec<(f64, Policy<Pk>)> {
+    debug_assert!(k <= policy_vec.len());
+
+    let mut ret: Vec<(f64, Policy<Pk>)> = vec![];
+    for i in 0..policy_vec.len() {
+        let mut policies: Vec<Policy<Pk>> = policy_vec.clone();
+        policies.remove(i);
+        ret.push((
+            prob / policy_vec.len() as f64,
+            Policy::<Pk>::Threshold(k, policies),
+        ));
     }
+    ret
 }