Merge #342: Tr compiler v2

ef4249d Add taproot compiler example usage (Aman Rojjha)
cef2a5b Add Tr-compiler write-up and doc-comment (Aman Rojjha)
5233c66 Add Taproot compiler API (Aman Rojjha)

Pull request description:

  This PR builds on top of #291. This aims to introduce an efficient version of the tapscript compiler by using a few heuristics to optimize over the expected average total cost for the TapTree.

  ## Strategy implemented

  - While merging TapTrees `A` and `B`, check whether the compilation of `Policy::Or(policy(A), policy(B))` is more efficient than a TapTree with the respective children `A` and `B`.

  **Note**: This doesn't include the `thresh(k, ...n..)` enumeration strategy. Planning on working on it separately.

ACKs for top commit:
  sanket1729:
    ACK ef4249d

Tree-SHA512: ae5949b5170ff4cc8202434655af41b4938e96801209896d117bf4e4bd85c0becc6fd0c7affb100c31655f0085c1ff5e7c19184db3433b2831f73db22d94348d

Author: sanket1729

Cargo.toml

@@ -28,6 +28,7 @@ hashbrown = { version = "0.11", optional = true }
 [dev-dependencies]
 bitcoind = {version = "0.26.1", features=["22_0"]}
 actual-rand = { package = "rand", version = "0.8.4"}
+secp256k1 = {version = "0.22.1", features = ["rand-std"]}
 
 [[example]]
 name = "htlc"
@@ -52,3 +53,7 @@ required-features = ["std"]
 [[example]]
 name = "xpub_descriptors"
 required-features = ["std"]
+
+[[example]]
+name = "taproot"
+required-features = ["compiler","std"]

contrib/test.sh

@@ -53,6 +53,7 @@ then
     cargo run --example verify_tx > /dev/null
     cargo run --example psbt
     cargo run --example xpub_descriptors
+    cargo run --example taproot --features=compiler
 fi
 
 if [ "$DO_NO_STD" = true ]

doc/Tr Compiler.pdf

@@ -0,0 +1,146 @@
+use std::collections::HashMap;
+use std::str::FromStr;
+
+use bitcoin::hashes::{hash160, sha256};
+use bitcoin::util::address::WitnessVersion;
+use bitcoin::Network;
+use miniscript::descriptor::DescriptorType;
+use miniscript::policy::Concrete;
+use miniscript::{Descriptor, Miniscript, Tap, TranslatePk, Translator};
+use secp256k1::{rand, KeyPair};
+
+// Refer to https://github.com/sanket1729/adv_btc_workshop/blob/master/workshop.md#creating-a-taproot-descriptor
+// for a detailed explanation of the policy and it's compilation
+
+struct StrPkTranslator {
+    pk_map: HashMap<String, bitcoin::XOnlyPublicKey>,
+}
+
+impl Translator<String, bitcoin::XOnlyPublicKey, ()> for StrPkTranslator {
+    fn pk(&mut self, pk: &String) -> Result<bitcoin::XOnlyPublicKey, ()> {
+        self.pk_map.get(pk).copied().ok_or(())
+    }
+
+    fn pkh(&mut self, _pkh: &String) -> Result<hash160::Hash, ()> {
+        unreachable!("Policy doesn't contain any pkh fragment");
+    }
+
+    fn sha256(&mut self, _sha256: &String) -> Result<sha256::Hash, ()> {
+        unreachable!("Policy does not contain any sha256 fragment");
+    }
+}
+
+fn main() {
+    let pol_str = "or(
+        99@thresh(2,
+            pk(hA), pk(S)
+        ),1@or(
+            99@pk(Ca),
+            1@and(pk(In), older(9))
+            )
+        )"
+    .replace(&[' ', '\n', '\t'][..], "");
+
+    let pol: Concrete<String> = Concrete::from_str(&pol_str).unwrap();
+    // In case we can't find an internal key for the given policy, we set the internal key to
+    // a random pubkey as specified by BIP341 (which are *unspendable* by any party :p)
+    let desc = pol.compile_tr(Some("UNSPENDABLE_KEY".to_string())).unwrap();
+
+    let expected_desc =
+        Descriptor::<String>::from_str("tr(Ca,{and_v(v:pk(In),older(9)),multi_a(2,hA,S)})")
+            .unwrap();
+    assert_eq!(desc, expected_desc);
+
+    // Check whether the descriptors are safe.
+    assert!(desc.sanity_check().is_ok());
+
+    // Descriptor Type and Version should match respectively for Taproot
+    let desc_type = desc.desc_type();
+    assert_eq!(desc_type, DescriptorType::Tr);
+    assert_eq!(desc_type.segwit_version().unwrap(), WitnessVersion::V1);
+
+    if let Descriptor::Tr(ref p) = desc {
+        // Check if internal key is correctly inferred as Ca
+        // assert_eq!(p.internal_key(), &pubkeys[2]);
+        assert_eq!(p.internal_key(), "Ca");
+
+        // Iterate through scripts
+        let mut iter = p.iter_scripts();
+        assert_eq!(
+            iter.next().unwrap(),
+            (
+                1u8,
+                &Miniscript::<String, Tap>::from_str("and_v(vc:pk_k(In),older(9))").unwrap()
+            )
+        );
+        assert_eq!(
+            iter.next().unwrap(),
+            (
+                1u8,
+                &Miniscript::<String, Tap>::from_str("multi_a(2,hA,S)").unwrap()
+            )
+        );
+        assert_eq!(iter.next(), None);
+    }
+
+    let mut pk_map = HashMap::new();
+
+    // We require secp for generating a random XOnlyPublicKey
+    let secp = secp256k1::Secp256k1::new();
+    let key_pair = KeyPair::new(&secp, &mut rand::thread_rng());
+    // Random unspendable XOnlyPublicKey provided for compilation to Taproot Descriptor
+    let unspendable_pubkey = bitcoin::XOnlyPublicKey::from_keypair(&key_pair);
+
+    pk_map.insert("UNSPENDABLE_KEY".to_string(), unspendable_pubkey);
+    let pubkeys = hardcoded_xonlypubkeys();
+    pk_map.insert("hA".to_string(), pubkeys[0]);
+    pk_map.insert("S".to_string(), pubkeys[1]);
+    pk_map.insert("Ca".to_string(), pubkeys[2]);
+    pk_map.insert("In".to_string(), pubkeys[3]);
+    let mut t = StrPkTranslator { pk_map };
+
+    let real_desc = desc.translate_pk(&mut t).unwrap();
+
+    // Max Satisfaction Weight for compilation, corresponding to the script-path spend
+    // `multi_a(2,PUBKEY_1,PUBKEY_2) at taptree depth 1, having
+    // Max Witness Size = scriptSig len + control_block size + varint(script_size) + script_size +
+    //                     varint(max satisfaction elements) + max satisfaction size
+    //                  = 4 + 65 + 1 + 70 + 1 + 132
+    let max_sat_wt = real_desc.max_satisfaction_weight().unwrap();
+    assert_eq!(max_sat_wt, 273);
+
+    // Compute the bitcoin address and check if it matches
+    let network = Network::Bitcoin;
+    let addr = real_desc.address(network).unwrap();
+    let expected_addr = bitcoin::Address::from_str(
+        "bc1pcc8ku64slu3wu04a6g376d2s8ck9y5alw5sus4zddvn8xgpdqw2swrghwx",
+    )
+    .unwrap();
+    assert_eq!(addr, expected_addr);
+}
+
+fn hardcoded_xonlypubkeys() -> Vec<bitcoin::XOnlyPublicKey> {
+    let serialized_keys: [[u8; 32]; 4] = [
+        [
+            22, 37, 41, 4, 57, 254, 191, 38, 14, 184, 200, 133, 111, 226, 145, 183, 245, 112, 100,
+            42, 69, 210, 146, 60, 179, 170, 174, 247, 231, 224, 221, 52,
+        ],
+        [
+            194, 16, 47, 19, 231, 1, 0, 143, 203, 11, 35, 148, 101, 75, 200, 15, 14, 54, 222, 208,
+            31, 205, 191, 215, 80, 69, 214, 126, 10, 124, 107, 154,
+        ],
+        [
+            202, 56, 167, 245, 51, 10, 193, 145, 213, 151, 66, 122, 208, 43, 10, 17, 17, 153, 170,
+            29, 89, 133, 223, 134, 220, 212, 166, 138, 2, 152, 122, 16,
+        ],
+        [
+            50, 23, 194, 4, 213, 55, 42, 210, 67, 101, 23, 3, 195, 228, 31, 70, 127, 79, 21, 188,
+            168, 39, 134, 58, 19, 181, 3, 63, 235, 103, 155, 213,
+        ],
+    ];
+    let mut keys: Vec<bitcoin::XOnlyPublicKey> = vec![];
+    for idx in 0..4 {
+        keys.push(bitcoin::XOnlyPublicKey::from_slice(&serialized_keys[idx][..]).unwrap());
+    }
+    keys
+}

examples/taproot.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> {
@@ -257,10 +244,14 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
     /// The policy tree constructed by root-level disjunctions over [`Or`][`Policy::Or`] and
     /// [`Thresh`][`Policy::Threshold`](1, ..) 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)))]`. Each policy in the vector is compiled into
-    /// the respective miniscripts. A Huffman Tree is created from this vector which optimizes over
-    /// the probabilitity of satisfaction for the respective branch in the TapTree.
+    /// 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)))]`. Each policy in the vector is compiled
+    /// into the respective miniscripts. A Huffman Tree is created from this vector which optimizes
+    /// over the probabilitity of satisfaction for the respective branch in the TapTree.
+    ///
+    /// Refer to [this link](https://gist.github.com/SarcasticNastik/9e70b2b43375aab3e78c51e09c288c89)
+    /// or [doc/Tr compiler.pdf] in the root of the repository to understand why such compilation
+    /// is also *cost-efficient*.
     // TODO: We might require other compile errors for Taproot.
     #[cfg(feature = "compiler")]
     pub fn compile_tr(&self, unspendable_key: Option<Pk>) -> Result<Descriptor<Pk>, Error> {
@@ -276,7 +267,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)