From 5233c6645581e6e5a5afee55946193210f16b7d4 Mon Sep 17 00:00:00 2001
From: Aman Rojjha <aman.rojjha@research.iiit.ac.in>
Date: Wed, 20 Apr 2022 01:02:51 +0530
Subject: [PATCH] Add Taproot compiler API

---
 src/policy/concrete.rs | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs
index ad30a7311e..5dce55cb88 100644
--- a/src/policy/concrete.rs
+++ b/src/policy/concrete.rs
@@ -194,19 +194,6 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
-    /// Compile [`Policy::Or`] and [`Policy::Threshold`] according to odds
-    #[cfg(feature = "compiler")]
-    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 internal_key from policy tree.
     #[cfg(feature = "compiler")]
     fn extract_key(self, unspendable_key: Option<Pk>) -> Result<(Pk, Policy<Pk>), Error> {
@@ -276,7 +263,21 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
                     internal_key,
                     match policy {
                         Policy::Trivial => None,
-                        policy => Some(policy.compile_tr_policy()?),
+                        policy => {
+                            let vec_policies: Vec<_> = policy.to_tapleaf_prob_vec(1.0);
+                            let mut leaf_compilations: Vec<(OrdF64, Miniscript<Pk, Tap>)> = vec![];
+                            for (prob, pol) in vec_policies {
+                                // policy corresponding to the key (replaced by unsatisfiable) is skipped
+                                if pol == Policy::Unsatisfiable {
+                                    continue;
+                                }
+                                let compilation = compiler::best_compilation::<Pk, Tap>(&pol)?;
+                                compilation.sanity_check()?;
+                                leaf_compilations.push((OrdF64(prob), compilation));
+                            }
+                            let taptree = with_huffman_tree::<Pk>(leaf_compilations)?;
+                            Some(taptree)
+                        }
                     },
                 )?;
                 Ok(tree)