feat(compile): give each disjunctive branch its own taproot leaf

CHANGELOG.md

@@ -22,6 +22,7 @@ page. See [DEVELOPMENT_CYCLE.md](DEVELOPMENT_CYCLE.md) for more details.
 - Updated dependencies to resolve RUSTSEC advisories in rustls, rustls-webpki, webpki-roots, quinn-proto, h2 and zstd-sys; bumped async-hwi to 0.0.35.
 - Fixed `create_tx`, `create_dns_tx`, `create_sp_tx` and `bump_fee` building transactions over Bitcoin Core's 400,000 wu (100,000 vB) standard weight limit, which no default node relays
 - Fixed `key generate` silently producing a 24-word mnemonic for any `--entropy` value other than 12; values other than 12 and 24 are now rejected
+- Changed `compile --type tr` to give each spending path of a disjunctive policy (`or`, `thresh(1, ..)`) its own taproot leaf, weighted by branch probability, so only the path that is used is revealed on-chain. The internal key is still the randomized unspendable key and `r` is still returned. Policies that exceed 1024 taproot leaves are rejected.
 
 ## [4.0.0]
 

src/utils/descriptors.rs

@@ -11,12 +11,15 @@ use bdk_wallet::{
         secp256k1::{PublicKey, Scalar, SecretKey},
     },
     miniscript::{
-        ForEachKey, FromStrKey, Legacy, MiniscriptKey, ScriptContext, Tap, descriptor::TapTree,
-        policy::Concrete,
+        ForEachKey, FromStrKey, Legacy, MiniscriptKey, ScriptContext, Tap,
+        descriptor::TapTree,
+        policy::{Concrete, compiler::CompilerError},
     },
 };
 #[cfg(any(feature = "hwi", feature = "compiler"))]
 use clap::ValueEnum;
+#[cfg(feature = "compiler")]
+use std::{cmp::Reverse, collections::BinaryHeap};
 use std::{str::FromStr, sync::Arc};
 
 use bdk_wallet::keys::DescriptorPublicKey;
@@ -340,7 +343,7 @@ pub fn compile_policy<Pk: FromStrKey>(
                 .map_err(|e| Error::Generic(format!("Failed to tweak NUMS key: {e}")))?;
             let (xonly_internal_key, _) = internal_key_point.x_only_public_key();
 
-            let tree = TapTree::Leaf(Arc::new(policy.compile()?));
+            let tree = compile_tap_tree(&policy)?;
             let internal_key = Pk::from_str(&xonly_internal_key.to_string())
                 .map_err(|e| Error::Generic(format!("Failed to parse internal key: {e}")))?;
             Descriptor::new_tr(internal_key, Some(tree))
@@ -349,6 +352,102 @@ pub fn compile_policy<Pk: FromStrKey>(
 
     Ok((descriptor.to_string(), r))
 }
+/// Upper bound on taproot leaves, matching rust-miniscript's `MAX_COMPILATION_LEAVES`.
+#[cfg(feature = "compiler")]
+const MAX_TAPLEAVES: usize = 1024;
+
+/// Probability weight with a total order, so it can key a `BinaryHeap`.
+#[cfg(feature = "compiler")]
+#[derive(PartialEq)]
+struct Weight(f64);
+
+#[cfg(feature = "compiler")]
+impl Eq for Weight {}
+
+#[cfg(feature = "compiler")]
+impl PartialOrd for Weight {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[cfg(feature = "compiler")]
+impl Ord for Weight {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.0.total_cmp(&other.0)
+    }
+}
+
+/// Flatten root-level disjunctions (`or`, `thresh(1, ..)`) into `(probability, branch)`
+/// pairs, one per spending path.
+#[cfg(feature = "compiler")]
+fn tapleaf_branches<Pk: MiniscriptKey>(
+    policy: &Concrete<Pk>,
+    prob: f64,
+    out: &mut Vec<(f64, Concrete<Pk>)>,
+) {
+    match policy {
+        Concrete::Or(subs) => {
+            let total: usize = subs.iter().map(|(weight, _)| *weight).sum();
+            for (weight, sub) in subs {
+                tapleaf_branches(sub, prob * *weight as f64 / total as f64, out);
+            }
+        }
+        Concrete::Thresh(thresh) if thresh.is_or() => {
+            for sub in thresh.iter() {
+                tapleaf_branches(sub, prob / thresh.n() as f64, out);
+            }
+        }
+        other => out.push((prob, other.clone())),
+    }
+}
+
+/// Compile `policy` into a taproot script tree with one leaf per spending path,
+/// Huffman-weighted by branch probability.
+///
+/// This stands in for rust-miniscript's `compile_tr`, which picks a key from the policy
+/// as the internal key. The caller needs to keep its own verifiable unspendable internal
+/// key, so only the tree is built here.
+#[cfg(feature = "compiler")]
+fn compile_tap_tree<Pk: MiniscriptKey>(policy: &Concrete<Pk>) -> Result<TapTree<Pk>, Error> {
+    // Validate the whole policy first, as `Concrete::compile` does. Checking each branch
+    // alone would miss problems that span branches, such as a key reused in two of them.
+    policy.is_valid().map_err(CompilerError::from)?;
+    match policy.is_safe_nonmalleable() {
+        (false, _) => return Err(CompilerError::TopLevelNonSafe.into()),
+        (_, false) => return Err(CompilerError::ImpossibleNonMalleableCompilation.into()),
+        _ => {}
+    }
+
+    let mut branches = Vec::new();
+    tapleaf_branches(policy, 1.0, &mut branches);
+
+    if branches.len() > MAX_TAPLEAVES {
+        return Err(Error::Generic(format!(
+            "Policy compiles to {} taproot leaves, the maximum is {MAX_TAPLEAVES}",
+            branches.len()
+        )));
+    }
+
+    let mut heap = BinaryHeap::new();
+    for (prob, branch) in branches {
+        let leaf: Miniscript<Pk, Tap> = branch.compile()?;
+        leaf.sanity_check()
+            .map_err(|e| Error::Generic(format!("Unsafe policy branch: {e}")))?;
+        heap.push((Reverse(Weight(prob)), TapTree::Leaf(Arc::new(leaf))));
+    }
+
+    while heap.len() > 1 {
+        let (Reverse(Weight(p1)), left) = heap.pop().expect("heap has two nodes");
+        let (Reverse(Weight(p2)), right) = heap.pop().expect("heap has two nodes");
+        heap.push((Reverse(Weight(p1 + p2)), TapTree::combine(left, right)));
+    }
+
+    heap.pop()
+        .map(|(_, tree)| tree)
+        .ok_or_else(|| Error::Generic("Policy has no spending path".to_string()))
+}
+
 /// Strip a trailing descriptor checksum.
 #[cfg(feature = "hwi")]
 pub(crate) fn strip_checksum(descriptor: &str) -> &str {
@@ -479,3 +578,69 @@ mod multipath_tests {
         assert_eq!(T::from_derivation_path(&p("m/84/1/0")), None);
     }
 }
+
+#[cfg(all(test, feature = "compiler"))]
+mod compile_tap_tree_tests {
+    use super::*;
+
+    /// `(depth, script)` of every leaf, in tree order.
+    fn leaves(policy: &str) -> Vec<(usize, String)> {
+        let policy: Concrete<String> = Concrete::from_str(policy).unwrap();
+        compile_tap_tree(&policy)
+            .unwrap()
+            .iter()
+            .map(|(depth, ms)| (usize::from(depth), ms.to_string()))
+            .collect()
+    }
+
+    #[test]
+    fn splits_or_into_one_leaf_per_branch() {
+        let found = leaves("or(pk(A),pk(B))");
+        assert_eq!(found.len(), 2);
+        assert!(found.iter().all(|(depth, _)| *depth == 1));
+    }
+
+    #[test]
+    fn flattens_nested_disjunctions() {
+        assert_eq!(leaves("or(pk(A),or(pk(B),pk(C)))").len(), 3);
+        assert_eq!(leaves("thresh(1,pk(A),pk(B),pk(C),pk(D))").len(), 4);
+    }
+
+    #[test]
+    fn puts_likely_branches_nearer_the_root() {
+        // A is 10x as likely as B or C, so the Huffman tree keeps it one step from the root.
+        let found = leaves("or(10@pk(A),or(1@pk(B),1@pk(C)))");
+        let depth_of = |script: &str| found.iter().find(|(_, s)| s == script).unwrap().0;
+
+        assert_eq!(depth_of("pk(A)"), 1);
+        assert_eq!(depth_of("pk(B)"), 2);
+        assert_eq!(depth_of("pk(C)"), 2);
+    }
+
+    #[test]
+    fn single_key_is_one_leaf_at_the_root() {
+        assert_eq!(leaves("pk(A)"), vec![(0, "pk(A)".to_string())]);
+    }
+
+    #[test]
+    fn keeps_a_conjunction_in_one_leaf() {
+        assert_eq!(leaves("and(pk(A),pk(B))").len(), 1);
+    }
+
+    #[test]
+    fn only_splits_one_of_n_thresholds() {
+        assert_eq!(leaves("thresh(2,pk(A),pk(B),pk(C))").len(), 1);
+    }
+
+    #[test]
+    fn splits_a_disjunction_below_a_conjunction_only_at_the_root() {
+        // The `or` sits under an `and`, so it is not a root-level disjunction.
+        assert_eq!(leaves("and(or(pk(A),pk(B)),pk(C))").len(), 1);
+    }
+
+    #[test]
+    fn same_policy_always_gives_the_same_tree() {
+        let policy = "thresh(1,pk(A),pk(B),pk(C),pk(D),pk(E))";
+        assert_eq!(leaves(policy), leaves(policy));
+    }
+}

tests/integration/init.rs

@@ -693,6 +693,178 @@ mod test_compile {
             .failure()
             .stderr(predicate::str::contains("Invalid policy"));
     }
+
+    const KEY_A: &str = "02e5b88fdb71c696e1a473f309a47535b7190e21a22bd25e7fc8bd055db3bba12f";
+    const KEY_B: &str = "03a34b99f22c790c4e36b2b3c2c35a36db06226e41c692fc82b8b56ac1c540c5bd";
+    const KEY_C: &str = "025476c2e83188368da1ff3e292e7acafcdb3566bb0ad253f62fc70f07aeee6357";
+
+    /// A disjunction compiled for `tr` gets one leaf per spending path, so only the path
+    /// that is used gets revealed on-chain.
+    #[test]
+    fn test_compile_tr_splits_or_into_separate_leaves() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("or(pk({KEY_A}),pk({KEY_B}))");
+
+        cli.cmd("compile", &[&policy, "--type", "tr"])
+            .assert()
+            .success()
+            .stdout(predicate::str::contains(format!(
+                ",{{pk({KEY_B}),pk({KEY_A})}}"
+            )))
+            .stdout(predicate::str::contains("or_b").not());
+    }
+
+    #[test]
+    fn test_compile_tr_splits_one_of_n_threshold_into_separate_leaves() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("thresh(1,pk({KEY_A}),pk({KEY_B}),pk({KEY_C}))");
+
+        cli.cmd("compile", &[&policy, "--type", "tr"])
+            .assert()
+            .success()
+            .stdout(predicate::str::contains(format!("pk({KEY_A})")))
+            .stdout(predicate::str::contains(format!("pk({KEY_B})")))
+            .stdout(predicate::str::contains(format!("pk({KEY_C})")))
+            .stdout(predicate::str::contains("multi_a").not());
+    }
+
+    /// The internal key is the randomized unspendable key and `r` is returned, never a key
+    /// taken from the policy.
+    #[test]
+    fn test_compile_tr_uses_unspendable_internal_key() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("or(pk({KEY_A}),pk({KEY_B}))");
+
+        let output = cli
+            .cmd("compile", &[&policy, "--type", "tr"])
+            .output()
+            .unwrap();
+        assert!(output.status.success());
+
+        let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap();
+        let descriptor = json["descriptor"].as_str().unwrap();
+        let internal_key = descriptor
+            .strip_prefix("tr(")
+            .and_then(|d| d.split(',').next())
+            .unwrap();
+
+        // The internal key is a 32-byte x-only key that is none of the policy keys.
+        assert_eq!(internal_key.len(), 64);
+        assert!(!KEY_A.contains(internal_key) && !KEY_B.contains(internal_key));
+        assert_eq!(json["r"].as_str().unwrap().len(), 64);
+    }
+
+    /// A single key used to fail under rust-miniscript's `compile_tr`.
+    #[test]
+    fn test_compile_tr_single_key_policy() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("pk({KEY_A})");
+
+        cli.cmd("compile", &[&policy, "--type", "tr"])
+            .assert()
+            .success()
+            .stdout(predicate::str::contains(format!("pk({KEY_A}))")));
+    }
+
+    /// A conjunction has a single spending path, so there is nothing to split.
+    #[test]
+    fn test_compile_tr_keeps_conjunction_in_one_leaf() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("and(pk({KEY_A}),pk({KEY_B}))");
+
+        cli.cmd("compile", &[&policy, "--type", "tr"])
+            .assert()
+            .success()
+            .stdout(predicate::str::contains("and_v("))
+            .stdout(predicate::str::contains(",{").not());
+    }
+
+    #[test]
+    fn test_compile_tr_keeps_k_of_n_threshold_in_one_leaf() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("thresh(2,pk({KEY_A}),pk({KEY_B}),pk({KEY_C}))");
+
+        cli.cmd("compile", &[&policy, "--type", "tr"])
+            .assert()
+            .success()
+            .stdout(predicate::str::contains("multi_a(2,"))
+            .stdout(predicate::str::contains(",{").not());
+    }
+
+    /// Validation runs on the whole policy, not per leaf: a key reused across two
+    /// branches is rejected just as it is for the other script types.
+    #[test]
+    fn test_compile_tr_rejects_key_reused_across_branches() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("or(pk({KEY_A}),and(pk({KEY_A}),pk({KEY_B})))");
+
+        for script_type in ["tr", "wsh"] {
+            cli.cmd("compile", &[&policy, "--type", script_type])
+                .assert()
+                .failure()
+                .stderr(predicate::str::contains("duplicate keys"));
+        }
+    }
+
+    #[test]
+    fn test_compile_tr_rejects_unsafe_branch() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        let policy = format!("or(pk({KEY_A}),older(10))");
+
+        cli.cmd("compile", &[&policy, "--type", "tr"])
+            .assert()
+            .failure()
+            .stderr(predicate::str::contains("not safe"));
+    }
+
+    /// Uses placeholder keys, since only the number of branches matters here.
+    fn one_of_n_policy(n: usize) -> String {
+        let keys: Vec<String> = (1..=n).map(|i| format!("pk(K{i:04})")).collect();
+        format!("thresh(1,{})", keys.join(","))
+    }
+
+    #[test]
+    fn test_compile_tr_accepts_the_maximum_number_of_leaves() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        cli.cmd(
+            "compile",
+            &[&one_of_n_policy(1024), "--type", "tr", "--placeholders"],
+        )
+        .assert()
+        .success();
+    }
+
+    #[test]
+    fn test_compile_tr_rejects_more_than_the_maximum_number_of_leaves() {
+        let temp_dir = TempDir::new().unwrap();
+        let cli = BdkCli::new("testnet", Some(temp_dir.path().to_path_buf()));
+
+        cli.cmd(
+            "compile",
+            &[&one_of_n_policy(1025), "--type", "tr", "--placeholders"],
+        )
+        .assert()
+        .failure()
+        .stderr(predicate::str::contains("taproot leaves"));
+    }
 }
 
 // --- CONFIG COMMAND TESTS ---