descriptor_tree/
lib.rs

1// SPDX-License-Identifier: CC0-1.0
2
3use miniscript::{
4    Legacy, Miniscript, MiniscriptKey, ScriptContext, Segwitv0, Threshold,
5    descriptor::{Descriptor, Sh, ShInner, SortedMultiVec, Tr, Wsh, WshInner},
6    miniscript::decode::Terminal,
7};
8
9type DescriptorTreeThreshold<Pk> = Threshold<DescriptorTree<Pk>, 0>;
10type KeylessDescriptorTreeThreshold<Pk> = Threshold<KeylessDescriptorTree<Pk>, 0>;
11
12/// A tree can be keyless, a key, or a threshold of trees
13#[derive(Clone, Debug)]
14pub enum DescriptorTree<Pk: MiniscriptKey> {
15    /// A keyless tree that can be either satisfiable or unsatisfiable
16    Keyless(bool),
17    /// A key
18    Key(Pk),
19    /// A threshold of trees
20    Threshold(DescriptorTreeThreshold<Pk>),
21}
22
23/// A tree can a key, or a threshold of trees
24#[derive(Clone, Debug)]
25pub enum KeylessDescriptorTree<Pk: MiniscriptKey> {
26    /// A key
27    Key(Pk),
28    /// A threshold of trees
29    Threshold(KeylessDescriptorTreeThreshold<Pk>),
30}
31
32/// A trait to construct a descriptor tree
33pub trait ToDescriptorTree<Pk: MiniscriptKey> {
34    /// Returns a descriptor tree
35    fn to_tree(&self) -> DescriptorTree<Pk>;
36}
37
38impl<Pk: MiniscriptKey> KeylessDescriptorTree<Pk> {
39    /// Returns a list of keys in the pruned descriptor
40    pub fn extract_keys(&self) -> Vec<Pk> {
41        match self {
42            KeylessDescriptorTree::Key(pk) => vec![pk.clone()],
43            KeylessDescriptorTree::Threshold(thresh) => {
44                thresh.iter().flat_map(|tree| tree.extract_keys()).collect()
45            }
46        }
47    }
48}
49
50impl<Pk: MiniscriptKey> DescriptorTree<Pk> {
51    /// Returns a list of keys in the descriptor
52    pub fn extract_keys(&self) -> Vec<Pk> {
53        match self {
54            DescriptorTree::Keyless(_) => Vec::new(),
55            DescriptorTree::Key(pk) => vec![pk.clone()],
56            DescriptorTree::Threshold(thresh) => {
57                thresh.iter().flat_map(|tree| tree.extract_keys()).collect()
58            }
59        }
60    }
61
62    /// Prune keyless leaves assuming satisifiable conditions are satisfied.
63    /// Sets new_k = max(old_k - num(satisfiable keyless leaves), 0) in each threshold.
64    pub fn prune_keyless(&self) -> Option<KeylessDescriptorTree<Pk>> {
65        let (_, pruned_tree) = self.prune_keyless_with_satisfiability();
66        pruned_tree
67    }
68
69    /// Returns pruned tree and whether its satisfiable
70    fn prune_keyless_with_satisfiability(&self) -> (bool, Option<KeylessDescriptorTree<Pk>>) {
71        match self {
72            DescriptorTree::Keyless(satisfiable) => (*satisfiable, None),
73            DescriptorTree::Key(pk) => (true, Some(KeylessDescriptorTree::Key(pk.clone()))),
74            DescriptorTree::Threshold(thresh) => {
75                let mut assume_satisfied = 0;
76                let mut keyed_subtrees = Vec::new();
77                for t in thresh.iter() {
78                    match t.prune_keyless_with_satisfiability() {
79                        (_, Some(subtree)) => {
80                            keyed_subtrees.push(subtree);
81                        }
82                        (true, None) => assume_satisfied += 1,
83                        (false, None) => {}
84                    }
85                }
86
87                let new_k = if thresh.k() > assume_satisfied {
88                    thresh.k() - assume_satisfied
89                } else {
90                    0
91                };
92
93                match (new_k, keyed_subtrees.len()) {
94                    (0, _) => (true, None),
95                    (1, 1) => (true, Some(keyed_subtrees.first().unwrap().clone())),
96                    (k, n) => {
97                        if k <= n {
98                            (
99                                true,
100                                KeylessDescriptorTreeThreshold::new(new_k, keyed_subtrees)
101                                    .ok()
102                                    .map(KeylessDescriptorTree::Threshold),
103                            )
104                        } else {
105                            (false, None)
106                        }
107                    }
108                }
109            }
110        }
111    }
112
113    fn from_ms_and<Ctx>(ms0: &Miniscript<Pk, Ctx>, ms1: &Miniscript<Pk, Ctx>) -> Self
114    where
115        Ctx: ScriptContext,
116    {
117        let tree0 = ms0.to_tree();
118        let tree1 = ms1.to_tree();
119        let thresh = DescriptorTreeThreshold::and(tree0, tree1);
120
121        DescriptorTree::Threshold(thresh)
122    }
123
124    fn from_ms_or<Ctx>(ms0: &Miniscript<Pk, Ctx>, ms1: &Miniscript<Pk, Ctx>) -> Self
125    where
126        Ctx: ScriptContext,
127    {
128        let tree0 = ms0.to_tree();
129        let tree1 = ms1.to_tree();
130        let thresh = DescriptorTreeThreshold::or(tree0, tree1);
131
132        DescriptorTree::Threshold(thresh)
133    }
134
135    fn from_sortedmulti<Ctx>(sortedmulti: &SortedMultiVec<Pk, Ctx>) -> Self
136    where
137        Ctx: ScriptContext,
138    {
139        let trees = sortedmulti
140            .pks()
141            .iter()
142            .map(|pk| DescriptorTree::Key(pk.clone()))
143            .collect();
144        let thresh = DescriptorTreeThreshold::new(sortedmulti.k(), trees).unwrap();
145
146        DescriptorTree::Threshold(thresh)
147    }
148}
149
150impl<Pk: MiniscriptKey> ToDescriptorTree<Pk> for Descriptor<Pk> {
151    fn to_tree(&self) -> DescriptorTree<Pk> {
152        match self {
153            Descriptor::Sh(sh) => sh.to_tree(),
154            Descriptor::Wsh(wsh) => wsh.to_tree(),
155            Descriptor::Tr(tr) => tr.to_tree(),
156            Descriptor::Wpkh(wpkh) => DescriptorTree::Key(wpkh.clone().into_inner()),
157            Descriptor::Pkh(pkh) => DescriptorTree::Key(pkh.clone().into_inner()),
158            Descriptor::Bare(bare) => bare.as_inner().to_tree(),
159        }
160    }
161}
162
163impl<Pk: MiniscriptKey> ToDescriptorTree<Pk> for Sh<Pk> {
164    fn to_tree(&self) -> DescriptorTree<Pk> {
165        match self.as_inner() {
166            ShInner::SortedMulti(sortedmulti) => {
167                DescriptorTree::from_sortedmulti::<Legacy>(sortedmulti)
168            }
169            ShInner::Wsh(wsh) => wsh.to_tree(),
170            ShInner::Wpkh(wpkh) => DescriptorTree::Key(wpkh.clone().into_inner()),
171            ShInner::Ms(ms) => ms.to_tree(),
172        }
173    }
174}
175
176impl<Pk: MiniscriptKey> ToDescriptorTree<Pk> for Wsh<Pk> {
177    fn to_tree(&self) -> DescriptorTree<Pk> {
178        match self.as_inner() {
179            WshInner::SortedMulti(sortedmulti) => {
180                DescriptorTree::from_sortedmulti::<Segwitv0>(sortedmulti)
181            }
182            WshInner::Ms(ms) => ms.to_tree(),
183        }
184    }
185}
186
187impl<Pk: MiniscriptKey> ToDescriptorTree<Pk> for Tr<Pk> {
188    fn to_tree(&self) -> DescriptorTree<Pk> {
189        let mut trees = Vec::new();
190        trees.push(DescriptorTree::Key(self.internal_key().clone()));
191
192        for (_, ms) in self.iter_scripts() {
193            trees.push(ms.to_tree());
194        }
195
196        let thresh = DescriptorTreeThreshold::or_n(trees);
197
198        DescriptorTree::Threshold(thresh)
199    }
200}
201
202impl<Pk: MiniscriptKey, Ctx: ScriptContext> ToDescriptorTree<Pk> for Miniscript<Pk, Ctx> {
203    fn to_tree(&self) -> DescriptorTree<Pk> {
204        match &self.node {
205            Terminal::True => DescriptorTree::Keyless(true),
206            Terminal::False => DescriptorTree::Keyless(false),
207            Terminal::PkK(pk) => DescriptorTree::Key(pk.clone()),
208            Terminal::PkH(pk) => DescriptorTree::Key(pk.clone()),
209            Terminal::RawPkH(_) => DescriptorTree::Keyless(true),
210            Terminal::After(_) => DescriptorTree::Keyless(true),
211            Terminal::Older(_) => DescriptorTree::Keyless(true),
212            Terminal::Sha256(_) => DescriptorTree::Keyless(true),
213            Terminal::Hash256(_) => DescriptorTree::Keyless(true),
214            Terminal::Ripemd160(_) => DescriptorTree::Keyless(true),
215            Terminal::Hash160(_) => DescriptorTree::Keyless(true),
216            Terminal::Alt(ms) => ms.to_tree(),
217            Terminal::Swap(ms) => ms.to_tree(),
218            Terminal::Check(ms) => ms.to_tree(),
219            Terminal::DupIf(ms) => ms.to_tree(),
220            Terminal::Verify(ms) => ms.to_tree(),
221            Terminal::NonZero(ms) => ms.to_tree(),
222            Terminal::ZeroNotEqual(ms) => ms.to_tree(),
223            Terminal::AndV(ms0, ms1) => DescriptorTree::from_ms_and(ms0, ms1),
224            Terminal::AndB(ms0, ms1) => DescriptorTree::from_ms_and(ms0, ms1),
225            Terminal::AndOr(ms0, ms1, ms2) => {
226                let and_tree = DescriptorTree::from_ms_and(ms0, ms1);
227                let or_tree = ms2.to_tree();
228                let thresh = DescriptorTreeThreshold::or(and_tree, or_tree);
229
230                DescriptorTree::Threshold(thresh)
231            }
232            Terminal::OrB(ms0, ms1) => DescriptorTree::from_ms_or(ms0, ms1),
233            Terminal::OrC(ms0, ms1) => DescriptorTree::from_ms_or(ms0, ms1),
234            Terminal::OrD(ms0, ms1) => DescriptorTree::from_ms_or(ms0, ms1),
235            Terminal::OrI(ms0, ms1) => DescriptorTree::from_ms_or(ms0, ms1),
236            Terminal::Thresh(thresh) => {
237                let mut trees = Vec::new();
238                for ms in thresh.iter() {
239                    let tree = ms.to_tree();
240                    trees.push(tree);
241                }
242                let thresh = DescriptorTreeThreshold::new(thresh.k(), trees).unwrap();
243
244                DescriptorTree::Threshold(thresh)
245            }
246            Terminal::Multi(thresh) => {
247                let trees = thresh
248                    .iter()
249                    .map(|pk| DescriptorTree::Key(pk.clone()))
250                    .collect();
251                let thresh = DescriptorTreeThreshold::new(thresh.k(), trees).unwrap();
252
253                DescriptorTree::Threshold(thresh)
254            }
255            Terminal::MultiA(thresh) => {
256                let trees = thresh
257                    .iter()
258                    .map(|pk| DescriptorTree::Key(pk.clone()))
259                    .collect();
260                let thresh = DescriptorTreeThreshold::new(thresh.k(), trees).unwrap();
261
262                DescriptorTree::Threshold(thresh)
263            }
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271    use bitcoin::secp256k1;
272    use miniscript::descriptor::DescriptorPublicKey;
273    use std::str::FromStr;
274
275    // Helper function to create test keys
276    fn create_test_key(index: u32) -> DescriptorPublicKey {
277        let secp = secp256k1::Secp256k1::new();
278        let secret_value = 1u32 + index;
279
280        let mut sk = [0u8; 32];
281        sk[28..32].copy_from_slice(&secret_value.to_be_bytes());
282
283        let pubkey = bitcoin::PublicKey {
284            inner: secp256k1::PublicKey::from_secret_key(
285                &secp,
286                &secp256k1::SecretKey::from_slice(&sk).expect("sk"),
287            ),
288            compressed: true,
289        };
290
291        DescriptorPublicKey::Single(miniscript::descriptor::SinglePub {
292            key: miniscript::descriptor::SinglePubKey::FullKey(pubkey),
293            origin: None,
294        })
295    }
296
297    #[test]
298    fn test_extract_keys_single() {
299        // Test with a single key
300        let key = create_test_key(1);
301        let tree = DescriptorTree::Key(key.clone());
302
303        let keys = tree.extract_keys();
304        assert_eq!(keys.len(), 1);
305        assert!(keys.contains(&key));
306    }
307
308    #[test]
309    fn test_extract_keys_threshold() {
310        // Create a 2-of-3 threshold
311        let key1 = create_test_key(1);
312        let key2 = create_test_key(2);
313        let key3 = create_test_key(3);
314
315        let trees = vec![
316            DescriptorTree::Key(key1.clone()),
317            DescriptorTree::Key(key2.clone()),
318            DescriptorTree::Key(key3.clone()),
319        ];
320
321        let thresh = DescriptorTreeThreshold::new(2, trees).unwrap();
322        let tree = DescriptorTree::Threshold(thresh);
323
324        let keys = tree.extract_keys();
325        assert_eq!(keys.len(), 3);
326        assert!(keys.contains(&key1));
327        assert!(keys.contains(&key2));
328        assert!(keys.contains(&key3));
329    }
330
331    #[test]
332    fn test_extract_keys_with_keyless() {
333        // Create a threshold with some keyless trees
334        let key1 = create_test_key(1);
335        let key2 = create_test_key(2);
336
337        let trees = vec![
338            DescriptorTree::Key(key1.clone()),
339            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
340            DescriptorTree::Key(key2.clone()),
341        ];
342
343        let thresh = DescriptorTreeThreshold::new(2, trees).unwrap();
344        let tree = DescriptorTree::Threshold(thresh);
345
346        let keys = tree.extract_keys();
347        assert_eq!(keys.len(), 2);
348        assert!(keys.contains(&key1));
349        assert!(keys.contains(&key2));
350    }
351
352    #[test]
353    fn test_prune_keyless_single() {
354        // Key trees should remain unchanged
355        let key = create_test_key(1);
356        let tree = DescriptorTree::Key(key.clone());
357
358        let result = tree.prune_keyless();
359        assert!(result.is_some());
360
361        match result.unwrap() {
362            KeylessDescriptorTree::Key(k) => assert_eq!(k, key),
363            _ => panic!("Expected Key tree"),
364        }
365
366        // Keyless trees should return None
367        let tree = DescriptorTree::Keyless::<DescriptorPublicKey>(true);
368        let result = tree.prune_keyless();
369        assert!(result.is_none());
370    }
371
372    #[test]
373    fn test_prune_keyless_threshold() {
374        // Create a 2-of-3 threshold with one keyless tree
375        let key1 = create_test_key(1);
376        let key2 = create_test_key(2);
377
378        let trees = vec![
379            DescriptorTree::Key(key1.clone()),
380            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
381            DescriptorTree::Keyless::<DescriptorPublicKey>(false),
382            DescriptorTree::Key(key2.clone()),
383        ];
384
385        let thresh = DescriptorTreeThreshold::new(2, trees).unwrap();
386        let tree = DescriptorTree::Threshold(thresh);
387
388        let result = tree.prune_keyless();
389        assert!(result.is_some());
390
391        match result.unwrap() {
392            KeylessDescriptorTree::Threshold(t) => {
393                assert_eq!(t.k(), 1);
394                assert_eq!(t.n(), 2);
395
396                // Check that the keys are the same
397                let mut keys = Vec::new();
398                for subtree in t.iter() {
399                    if let KeylessDescriptorTree::Key(pk) = subtree {
400                        keys.push(pk.clone())
401                    }
402                }
403                assert!(keys.contains(&key1));
404                assert!(keys.contains(&key2));
405            }
406            _ => panic!("Expected Threshold tree"),
407        }
408    }
409
410    #[test]
411    fn test_prune_keyless_reduce_threshold() {
412        // Test that k gets reduced to n if needed
413        let key1 = create_test_key(1);
414        let key2 = create_test_key(2);
415
416        let trees = vec![
417            DescriptorTree::Key(key1.clone()),
418            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
419            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
420            DescriptorTree::Keyless::<DescriptorPublicKey>(false),
421            DescriptorTree::Key(key2.clone()),
422        ];
423
424        let thresh = DescriptorTreeThreshold::new(3, trees).unwrap();
425        let tree = DescriptorTree::Threshold(thresh);
426
427        let result = tree.prune_keyless();
428        assert!(result.is_some());
429
430        match result.unwrap() {
431            KeylessDescriptorTree::Threshold(t) => {
432                assert_eq!(t.k(), 1); // Reduced from 3 to 1
433                assert_eq!(t.n(), 2);
434
435                // Check that the keys are the same
436                let mut keys = Vec::new();
437                for subtree in t.iter() {
438                    if let KeylessDescriptorTree::Key(pk) = subtree {
439                        keys.push(pk.clone())
440                    }
441                }
442                assert!(keys.contains(&key1));
443                assert!(keys.contains(&key2));
444            }
445            _ => panic!("Expected Threshold tree"),
446        }
447    }
448
449    #[test]
450    fn test_prune_keyless_all_trees_keyless() {
451        // Test a scenario where all trees are keyless
452        let trees = vec![
453            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
454            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
455        ];
456
457        let thresh = DescriptorTreeThreshold::new(1, trees).unwrap();
458        let tree = DescriptorTree::Threshold(thresh);
459
460        let result = tree.prune_keyless();
461        assert!(result.is_none());
462    }
463
464    #[test]
465    fn test_prune_keyless_k_exceeds_n_due_to_unsatisfiability() {
466        let key1 = create_test_key(1);
467
468        // Test tree with satisfiable pruned leaf
469        let trees = vec![
470            DescriptorTree::Key(key1.clone()),
471            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
472        ];
473
474        let thresh = DescriptorTreeThreshold::new(2, trees).unwrap();
475        let tree = DescriptorTree::Threshold(thresh);
476
477        let result = tree.prune_keyless();
478        assert!(result.is_some());
479
480        // Test tree with unsatisfiable pruned leaf
481        let trees = vec![
482            DescriptorTree::Key(key1.clone()),
483            DescriptorTree::Keyless::<DescriptorPublicKey>(false),
484        ];
485
486        let thresh = DescriptorTreeThreshold::new(2, trees).unwrap();
487        let tree = DescriptorTree::Threshold(thresh);
488
489        let result = tree.prune_keyless();
490        assert!(result.is_none());
491    }
492
493    #[test]
494    fn test_prune_keyless_single_key_after_pruning() {
495        // Create a threshold with one key and one keyless tree
496        let key = create_test_key(1);
497
498        let trees = vec![
499            DescriptorTree::Key(key.clone()),
500            DescriptorTree::Keyless::<DescriptorPublicKey>(true),
501        ];
502
503        let thresh = DescriptorTreeThreshold::new(2, trees).unwrap();
504        let tree = DescriptorTree::Threshold(thresh);
505
506        let result = tree.prune_keyless();
507        assert!(result.is_some());
508
509        // After pruning, we should be left with just the key tree
510        match result.unwrap() {
511            KeylessDescriptorTree::Key(k) => assert_eq!(k, key),
512            _ => panic!("Expected a single Key tree"),
513        }
514    }
515
516    #[test]
517    fn test_to_tree_wpkh() {
518        // Test WPKH descriptor
519        let key = create_test_key(1);
520        let desc_str = format!("wpkh({})", key);
521        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
522
523        let tree = desc.to_tree();
524
525        // Should be a single key tree
526        match tree {
527            DescriptorTree::Key(k) => match (k, key) {
528                (
529                    DescriptorPublicKey::Single(single_k),
530                    DescriptorPublicKey::Single(single_key),
531                ) => match (&single_k.key, &single_key.key) {
532                    (
533                        miniscript::descriptor::SinglePubKey::FullKey(k_pk),
534                        miniscript::descriptor::SinglePubKey::FullKey(key_pk),
535                    ) => {
536                        assert_eq!(k_pk.inner.serialize(), key_pk.inner.serialize());
537                    }
538                    _ => panic!("Expected FullKey for both keys"),
539                },
540                _ => panic!("Expected Single keys"),
541            },
542            _ => panic!("Expected Key tree for wpkh descriptor"),
543        }
544    }
545
546    #[test]
547    fn test_to_tree_pkh() {
548        // Test PKH descriptor
549        let key = create_test_key(1);
550        let desc_str = format!("pkh({})", key);
551        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
552
553        let tree = desc.to_tree();
554
555        // Should be a single key tree
556        match tree {
557            DescriptorTree::Key(_) => {}
558            _ => panic!("Expected Key tree for pkh descriptor"),
559        }
560    }
561
562    #[test]
563    fn test_to_tree_bare_pk() {
564        let desc_str = "pk(020000000000000000000000000000000000000000000000000000000000000002)";
565        let desc = Descriptor::<DescriptorPublicKey>::from_str(desc_str).unwrap();
566
567        let tree = desc.to_tree();
568
569        // Should be a single key tree
570        match tree {
571            DescriptorTree::Key(_) => {}
572            _ => panic!("Expected Key tree for pkh descriptor"),
573        }
574    }
575
576    #[test]
577    fn test_to_tree_wsh_multi() {
578        // Test WSH with a multisig
579        let key1 = create_test_key(1);
580        let key2 = create_test_key(2);
581        let key3 = create_test_key(3);
582
583        let desc_str = format!("wsh(multi(2,{},{},{}))", key1, key2, key3);
584        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
585
586        let tree = desc.to_tree();
587
588        let keys = tree.extract_keys();
589        assert_eq!(keys.len(), 3);
590
591        // Should be a threshold with 3 key trees
592        match tree {
593            DescriptorTree::Threshold(t) => {
594                assert_eq!(t.k(), 2);
595                assert_eq!(t.n(), 3);
596
597                // We can't use contains directly because we need to compare the serialized keys
598                let key_serialized1 = serialize_descriptor_pubkey(&key1);
599                let key_serialized2 = serialize_descriptor_pubkey(&key2);
600                let key_serialized3 = serialize_descriptor_pubkey(&key3);
601
602                let extracted_serialized: Vec<_> =
603                    keys.iter().map(serialize_descriptor_pubkey).collect();
604
605                assert!(extracted_serialized.contains(&key_serialized1));
606                assert!(extracted_serialized.contains(&key_serialized2));
607                assert!(extracted_serialized.contains(&key_serialized3));
608            }
609            _ => panic!("Expected Threshold tree for wsh(multi) descriptor"),
610        }
611    }
612
613    #[test]
614    fn test_to_tree_sh_wsh_multi() {
615        // Test nested SH-WSH with a multisig
616        let key1 = create_test_key(1);
617        let key2 = create_test_key(2);
618
619        let desc_str = format!("sh(wsh(multi(2,{},{})))", key1, key2);
620        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
621
622        let tree = desc.to_tree();
623
624        let keys = tree.extract_keys();
625        assert_eq!(keys.len(), 2);
626
627        // Should be a threshold with 2 key trees
628        match tree {
629            DescriptorTree::Threshold(t) => {
630                assert_eq!(t.k(), 2);
631                assert_eq!(t.n(), 2);
632
633                let key_serialized1 = serialize_descriptor_pubkey(&key1);
634                let key_serialized2 = serialize_descriptor_pubkey(&key2);
635
636                let extracted_serialized: Vec<_> =
637                    keys.iter().map(serialize_descriptor_pubkey).collect();
638
639                assert!(extracted_serialized.contains(&key_serialized1));
640                assert!(extracted_serialized.contains(&key_serialized2));
641            }
642            _ => panic!("Expected Threshold tree for sh(wsh(multi)) descriptor"),
643        }
644    }
645
646    #[test]
647    fn test_to_tree_tr() {
648        // Test taproot descriptor with an internal key and no scripts
649        let key = create_test_key(1);
650        let desc_str = format!("tr({})", key);
651        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
652
653        let tree = desc.to_tree();
654
655        let keys = tree.extract_keys();
656        assert_eq!(keys.len(), 1);
657
658        // Should be a threshold tree with just the internal key
659        match tree {
660            DescriptorTree::Threshold(t) => {
661                assert_eq!(t.n(), 1); // Just the internal key, no scripts
662
663                // Compare serialized keys
664                let key_serialized = serialize_descriptor_pubkey(&key);
665                let extracted_serialized = serialize_descriptor_pubkey(&keys[0]);
666                assert_eq!(extracted_serialized, key_serialized);
667            }
668            _ => panic!("Expected Threshold tree for tr descriptor"),
669        }
670    }
671
672    #[test]
673    fn test_to_tree_tr_with_scripts() {
674        // Test taproot descriptor with scripts
675        let key1 = create_test_key(1);
676        let key2 = create_test_key(2);
677        let key3 = create_test_key(3);
678
679        // tr with internal key and one script path spending with key2
680        let desc_str = format!("tr({},pk({}))", key1, key2);
681        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
682
683        let tree = desc.to_tree();
684
685        let keys = tree.extract_keys();
686        assert_eq!(keys.len(), 2);
687
688        // Should be a threshold with internal key and script key
689        match tree {
690            DescriptorTree::Threshold(t) => {
691                assert_eq!(t.n(), 2); // Internal key + 1 script
692
693                let key_serialized1 = serialize_descriptor_pubkey(&key1);
694                let key_serialized2 = serialize_descriptor_pubkey(&key2);
695
696                let extracted_serialized: Vec<_> =
697                    keys.iter().map(serialize_descriptor_pubkey).collect();
698
699                assert!(extracted_serialized.contains(&key_serialized1));
700                assert!(extracted_serialized.contains(&key_serialized2));
701            }
702            _ => panic!("Expected Threshold tree for tr descriptor with script"),
703        }
704
705        // tr with internal key and two script paths
706        let desc_str = format!("tr({},{{pk({}),pk({})}})", key1, key2, key3);
707        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
708
709        let tree = desc.to_tree();
710
711        let keys = tree.extract_keys();
712        assert_eq!(keys.len(), 3);
713
714        // Should be a threshold with internal key and two script keys
715        match tree {
716            DescriptorTree::Threshold(t) => {
717                assert_eq!(t.n(), 3); // Internal key + 2 scripts
718
719                let key_serialized1 = serialize_descriptor_pubkey(&key1);
720                let key_serialized2 = serialize_descriptor_pubkey(&key2);
721                let key_serialized3 = serialize_descriptor_pubkey(&key3);
722
723                let extracted_serialized: Vec<_> =
724                    keys.iter().map(serialize_descriptor_pubkey).collect();
725
726                assert!(extracted_serialized.contains(&key_serialized1));
727                assert!(extracted_serialized.contains(&key_serialized2));
728                assert!(extracted_serialized.contains(&key_serialized3));
729            }
730            _ => panic!("Expected Threshold tree for tr descriptor with multiple scripts"),
731        }
732    }
733
734    #[test]
735    fn test_to_tree_with_miniscript_and() {
736        // Test a descriptor with an AND operation in miniscript
737        let key1 = create_test_key(1);
738        let key2 = create_test_key(2);
739
740        // wsh(and_v(v:pk(key1),pk(key2))) - requires both keys
741        let desc_str = format!("wsh(and_v(v:pk({}),pk({})))", key1, key2);
742        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
743
744        let tree = desc.to_tree();
745
746        let keys = tree.extract_keys();
747        assert_eq!(keys.len(), 2);
748
749        // Should be a threshold requiring all (k=n)
750        match tree {
751            DescriptorTree::Threshold(t) => {
752                assert_eq!(t.k(), t.n()); // AND requires all keys
753
754                let key_serialized1 = serialize_descriptor_pubkey(&key1);
755                let key_serialized2 = serialize_descriptor_pubkey(&key2);
756
757                let extracted_serialized: Vec<_> =
758                    keys.iter().map(serialize_descriptor_pubkey).collect();
759
760                assert!(extracted_serialized.contains(&key_serialized1));
761                assert!(extracted_serialized.contains(&key_serialized2));
762            }
763            _ => panic!("Expected Threshold tree for AND miniscript"),
764        }
765    }
766
767    #[test]
768    fn test_to_tree_with_miniscript_or() {
769        // Test a descriptor with an OR operation in miniscript
770        let key1 = create_test_key(1);
771        let key2 = create_test_key(2);
772
773        // wsh(or_d(pk(key1),pk(key2))) - requires one of the keys
774        let desc_str = format!("wsh(or_d(pk({}),pk({})))", key1, key2);
775        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
776
777        let tree = desc.to_tree();
778
779        let keys = tree.extract_keys();
780        assert_eq!(keys.len(), 2);
781
782        // Should be a 1-of-n threshold
783        match tree {
784            DescriptorTree::Threshold(t) => {
785                assert_eq!(t.k(), 1); // OR requires 1 of the keys
786                assert_eq!(t.n(), 2);
787
788                let key_serialized1 = serialize_descriptor_pubkey(&key1);
789                let key_serialized2 = serialize_descriptor_pubkey(&key2);
790
791                let extracted_serialized: Vec<_> =
792                    keys.iter().map(serialize_descriptor_pubkey).collect();
793
794                assert!(extracted_serialized.contains(&key_serialized1));
795                assert!(extracted_serialized.contains(&key_serialized2));
796            }
797            _ => panic!("Expected Threshold tree for OR miniscript"),
798        }
799    }
800
801    #[test]
802    fn test_to_tree_with_miniscript_thresh() {
803        // Test a descriptor with a threshold operation in miniscript
804        let key1 = create_test_key(1);
805        let key2 = create_test_key(2);
806        let key3 = create_test_key(3);
807
808        // wsh(thresh(2,pk(key1),pk(key2),pk(key3)))
809        let desc_str = format!("wsh(thresh(2,pk({}),s:pk({}),s:pk({})))", key1, key2, key3);
810        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
811
812        let tree = desc.to_tree();
813
814        let keys = tree.extract_keys();
815        assert_eq!(keys.len(), 3);
816
817        // Should be a 2-of-3 threshold
818        match tree {
819            DescriptorTree::Threshold(t) => {
820                assert_eq!(t.k(), 2);
821                assert_eq!(t.n(), 3);
822
823                let key_serialized1 = serialize_descriptor_pubkey(&key1);
824                let key_serialized2 = serialize_descriptor_pubkey(&key2);
825                let key_serialized3 = serialize_descriptor_pubkey(&key3);
826
827                let extracted_serialized: Vec<_> =
828                    keys.iter().map(serialize_descriptor_pubkey).collect();
829
830                assert!(extracted_serialized.contains(&key_serialized1));
831                assert!(extracted_serialized.contains(&key_serialized2));
832                assert!(extracted_serialized.contains(&key_serialized3));
833            }
834            _ => panic!("Expected Threshold tree for thresh miniscript"),
835        }
836    }
837
838    #[test]
839    fn test_to_tree_with_miniscript_andor() {
840        // Test a descriptor with an AND_OR operation in miniscript
841        let key1 = create_test_key(1);
842        let key2 = create_test_key(2);
843        let key3 = create_test_key(3);
844
845        // wsh(andor(pk(key1),pk(key2),pk(key3)))
846        // Equivalent to: (key1 AND key2) OR key3
847        let desc_str = format!("wsh(andor(pk({}),pk({}),pk({})))", key1, key2, key3);
848        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
849
850        let tree = desc.to_tree();
851
852        // Extract keys and ensure all 3 keys are there
853        let keys = tree.extract_keys();
854        assert_eq!(keys.len(), 3);
855
856        let key_serialized1 = serialize_descriptor_pubkey(&key1);
857        let key_serialized2 = serialize_descriptor_pubkey(&key2);
858        let key_serialized3 = serialize_descriptor_pubkey(&key3);
859
860        let extracted_serialized: Vec<_> = keys.iter().map(serialize_descriptor_pubkey).collect();
861
862        assert!(extracted_serialized.contains(&key_serialized1));
863        assert!(extracted_serialized.contains(&key_serialized2));
864        assert!(extracted_serialized.contains(&key_serialized3));
865    }
866
867    #[test]
868    fn test_to_tree_with_timelock() {
869        // Test a descriptor with a timelock
870        let key = create_test_key(1);
871
872        // wsh(and_v(v:pk(key),after(100)))
873        let desc_str = format!("wsh(and_v(v:pk({}),after(100)))", key);
874        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
875
876        let tree = desc.to_tree();
877
878        let keys = tree.extract_keys();
879        assert_eq!(keys.len(), 1);
880
881        // Should be a threshold with one real key and the timelock becomes keyless
882        match tree {
883            DescriptorTree::Threshold(_) => {
884                let key_serialized = serialize_descriptor_pubkey(&key);
885                let extracted_serialized = serialize_descriptor_pubkey(&keys[0]);
886                assert_eq!(extracted_serialized, key_serialized);
887
888                // Test prune_keyless behavior
889                let keyless = tree.prune_keyless();
890                assert!(keyless.is_some());
891
892                // After removing keyless, should just have a single key
893                match keyless.unwrap() {
894                    KeylessDescriptorTree::Key(k) => {
895                        let k_serialized = serialize_descriptor_pubkey(&k);
896                        assert_eq!(k_serialized, key_serialized);
897                    }
898                    _ => panic!("Expected single key after prune_keyless"),
899                }
900            }
901            _ => panic!("Expected Threshold tree for timelock descriptor"),
902        }
903    }
904
905    #[test]
906    fn test_to_tree_with_hashlocks() {
907        // Test a descriptor with hash locks
908        let key1 = create_test_key(1);
909        let key2 = create_test_key(2);
910
911        // wsh(or_d(pk(key1),and_v(v:pk(key2),sha256(7924b373d3b1a8269050c572a0b5a5461d8211d5777744c78ca247bc30569b21))))
912        let desc_str = format!(
913            "wsh(or_d(pk({}),and_v(v:pk({}),sha256(7924b373d3b1a8269050c572a0b5a5461d8211d5777744c78ca247bc30569b21))))",
914            key1, key2
915        );
916        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
917
918        let tree = desc.to_tree();
919
920        // Extract keys
921        let keys = tree.extract_keys();
922        assert_eq!(keys.len(), 2);
923
924        let key_serialized1 = serialize_descriptor_pubkey(&key1);
925        let key_serialized2 = serialize_descriptor_pubkey(&key2);
926
927        let extracted_serialized: Vec<_> = keys.iter().map(serialize_descriptor_pubkey).collect();
928
929        assert!(extracted_serialized.contains(&key_serialized1));
930        assert!(extracted_serialized.contains(&key_serialized2));
931
932        // Test prune_keyless
933        let keyless = tree.prune_keyless();
934        assert!(keyless.is_some());
935    }
936
937    #[test]
938    fn test_complex_descriptor() {
939        // Test a more complex descriptor with multiple conditions
940        let key1 = create_test_key(1);
941        let key2 = create_test_key(2);
942        let key3 = create_test_key(3);
943
944        // wsh(or_d(pk(key1),and_v(v:thresh(2,pk(key2),pk(key3),older(1000)),after(100))))
945        let desc_str = format!(
946            "wsh(or_d(pk({}),and_v(v:thresh(2,pk({}),s:pk({}),sln:older(1000)),after(100))))",
947            key1, key2, key3
948        );
949        let desc = Descriptor::<DescriptorPublicKey>::from_str(&desc_str).unwrap();
950
951        let tree = desc.to_tree();
952
953        // Extract keys
954        let keys = tree.extract_keys();
955        assert_eq!(keys.len(), 3);
956
957        let key_serialized1 = serialize_descriptor_pubkey(&key1);
958        let key_serialized2 = serialize_descriptor_pubkey(&key2);
959        let key_serialized3 = serialize_descriptor_pubkey(&key3);
960
961        let extracted_serialized: Vec<_> = keys.iter().map(serialize_descriptor_pubkey).collect();
962
963        assert!(extracted_serialized.contains(&key_serialized1));
964        assert!(extracted_serialized.contains(&key_serialized2));
965        assert!(extracted_serialized.contains(&key_serialized3));
966    }
967
968    // Helper function to serialize a descriptor public key for comparison
969    fn serialize_descriptor_pubkey(key: &DescriptorPublicKey) -> Vec<u8> {
970        match key {
971            DescriptorPublicKey::Single(single) => match &single.key {
972                miniscript::descriptor::SinglePubKey::FullKey(pk) => pk.inner.serialize().to_vec(),
973                miniscript::descriptor::SinglePubKey::XOnly(xpk) => xpk.serialize().to_vec(),
974            },
975            DescriptorPublicKey::XPub(xpub) => xpub.xkey.encode().to_vec(),
976            DescriptorPublicKey::MultiXPub(multi) => multi.xkey.encode().to_vec(),
977        }
978    }
979}
descriptor_tree/lib.rs

descriptor_tree/
lib.rs
