Implement tr-compiler-v3 enumeration compiler

SarcasticNastik · SarcasticNastik · commit 788a03b85e9e · 2022-06-27T23:56:00.000+05:30
diff --git a/src/lib.rs b/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,
diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs
@@ -187,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_COMPILATION_LEAVES`].
     #[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()
@@ -199,17 +204,138 @@ 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())],
+            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_COMPILATION_LEAVES || 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> {
@@ -306,6 +432,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
@@ -944,3 +1120,30 @@ fn with_huffman_tree<Pk: MiniscriptKey>(
         .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 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
+}
diff --git a/src/policy/mod.rs b/src/policy/mod.rs
@@ -487,4 +487,118 @@ mod tests {
             assert_eq!(descriptor, expected_descriptor);
         }
     }
+
+    #[test]
+    #[cfg(feature = "compiler")]
+    fn private_enumerative_compiler() {
+        // 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_private_experimental(Some(unspendable_key.clone()))
+                .unwrap();
+            let ms_compilations: [Miniscript<String, Tap>; 6] = [
+                ms_str!("multi_a(2,A,C)"),
+                ms_str!("multi_a(2,A,B)"),
+                ms_str!("multi_a(2,C,D)"),
+                ms_str!("multi_a(2,B,D)"),
+                ms_str!("multi_a(2,B,C)"),
+                ms_str!("multi_a(2,A,D)"),
+            ];
+
+            let tree = TapTree::Tree(
+                Arc::new(TapTree::Tree(
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[0].clone()))),
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[1].clone()))),
+                )),
+                Arc::new(TapTree::Tree(
+                    Arc::new(TapTree::Tree(
+                        Arc::new(TapTree::Leaf(Arc::new(ms_compilations[2].clone()))),
+                        Arc::new(TapTree::Leaf(Arc::new(ms_compilations[3].clone()))),
+                    )),
+                    Arc::new(TapTree::Tree(
+                        Arc::new(TapTree::Leaf(Arc::new(ms_compilations[4].clone()))),
+                        Arc::new(TapTree::Leaf(Arc::new(ms_compilations[5].clone()))),
+                    )),
+                )),
+            );
+
+            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(or(and(pk(A),pk(B)),and(pk(E),pk(F))),and(pk(C),pk(D)))");
+            let descriptor = policy
+                .compile_tr_private_experimental(Some(unspendable_key.clone()))
+                .unwrap();
+
+            let ms_compilations: [Miniscript<String, Tap>; 3] = [
+                ms_str!("and_v(v:pk(A),pk(B))"),
+                ms_str!("and_v(v:pk(E),pk(F))"),
+                ms_str!("and_v(v:pk(C),pk(D))"),
+            ];
+            let tree = TapTree::Tree(
+                Arc::new(TapTree::Leaf(Arc::new(ms_compilations[2].clone()))),
+                Arc::new(TapTree::Tree(
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[1].clone()))),
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[0].clone()))),
+                )),
+            );
+            let expected_descriptor =
+                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
+            assert_eq!(descriptor, expected_descriptor);
+        }
+
+        {
+            let policy: Concrete<String> = policy_str!(
+                "or(or(and(pk(A),pk(B)),and(pk(E),pk(F))),thresh(1,and(pk(C),pk(D)),and(pk(G),pk(H))))"
+            );
+            let descriptor = policy
+                .compile_tr_private_experimental(Some(unspendable_key.clone()))
+                .unwrap();
+
+            let ms_compilations: [Miniscript<String, Tap>; 4] = [
+                ms_str!("and_v(v:pk(A),pk(B))"),
+                ms_str!("and_v(v:pk(C),pk(D))"),
+                ms_str!("and_v(v:pk(E),pk(F))"),
+                ms_str!("and_v(v:pk(G),pk(H))"),
+            ];
+            let tree = TapTree::Tree(
+                Arc::new(TapTree::Tree(
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[3].clone()))),
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[2].clone()))),
+                )),
+                Arc::new(TapTree::Tree(
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[1].clone()))),
+                    Arc::new(TapTree::Leaf(Arc::new(ms_compilations[0].clone()))),
+                )),
+            );
+            let expected_descriptor =
+                Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap();
+            assert_eq!(descriptor, expected_descriptor);
+        }
+
+        {
+            // Invalid policy compilation (Duplicate PubKeys)
+            let policy: Concrete<String> = policy_str!("or(and(pk(A),pk(B)),and(pk(A),pk(D)))");
+            let descriptor = policy.compile_tr_private_experimental(Some(unspendable_key.clone()));
+
+            assert_eq!(
+                descriptor.unwrap_err().to_string(),
+                "Policy contains duplicate keys"
+            );
+        }
+
+        {
+            let policy: Concrete<String> = policy_str!("thresh(51,pk(A),pk(B),pk(C),pk(D),pk(E),pk(F),pk(G),pk(H),pk(I),pk(J),pk(K),pk(L),pk(M),pk(N),pk(O),pk(P),pk(Q),pk(R),pk(S),pk(T),pk(U),pk(V),pk(W),pk(X),pk(Y),pk(Z),pk(AA),pk(BB),pk(CC),pk(DD),pk(EE),pk(FF),pk(GG),pk(HH),pk(II),pk(JJ),pk(KK),pk(LL),pk(MM),pk(NN),pk(OO),pk(PP),pk(QQ),pk(RR),pk(SS),pk(TT),pk(UU),pk(VV),pk(WW),pk(XX),pk(YY),pk(ZZ))");
+            let descriptor = dbg!(policy
+                .compile_tr_private_experimental(Some(unspendable_key.clone()))
+                .unwrap());
+        }
+    }
 }
