bitcoin_scripts/
taproot.rs

1// BP foundation libraries Bitcoin crates implementing the foundations of
2// Bitcoin protocol by LNP/BP Association (https://lnp-bp.org)
3//
4// Written in 2020-2022 by
5//     Dr. Maxim Orlovsky <orlovsky@lnp-bp.org>
6//
7// This software is distributed without any warranty.
8//
9// You should have received a copy of the Apache-2.0 License
10// along with this software.
11// If not, see <https://opensource.org/licenses/Apache-2.0>.
12
13//! Taproot script tree implementation allowing arbitrary tree processing/
14//! modification (see [`TaprootScriptTree`] structure).
15
16use std::borrow::{Borrow, BorrowMut};
17use std::cmp::Ordering;
18use std::fmt::{self, Debug, Display, Formatter};
19use std::io::{Read, Write};
20use std::ops::{Deref, Not};
21use std::str::FromStr;
22
23use amplify::Wrapper;
24use bitcoin::hashes::Hash;
25use bitcoin::psbt::TapTree;
26use bitcoin::util::taproot::{LeafVersion, TapBranchHash, TapLeafHash, TaprootBuilder};
27use bitcoin::Script;
28use strict_encoding::{StrictDecode, StrictEncode};
29
30use crate::types::IntoNodeHash;
31use crate::{LeafScript, TapNodeHash, TapScript};
32
33/// Taproot tree or subtree construction error: improper lexicographi ordering
34/// of the nodes.
35#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Error, Display)]
36#[display(
37    "invalid taproot tree lexicographic node ordering in branch {dfs_path}, where the hash of the \
38     left-side child {left_hash} is larger than the hash of the right-side child {right_hash}"
39)]
40pub struct TaprootTreeError {
41    /// Node hash of the left-side child node.
42    pub left_hash: TapNodeHash,
43    /// Node hash of the right-side child node.
44    pub right_hash: TapNodeHash,
45    /// Path of the node in DFS (depth-first search) order.
46    pub dfs_path: DfsPath,
47}
48
49/// Error indicating that the maximum taproot script tree depth exceeded.
50#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Error, Display)]
51#[display("maximum taproot script tree depth exceeded.")]
52pub struct MaxDepthExceeded;
53
54/// Error indicating an attempt to raise subtree above its depth (i.e. root).
55#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Error, Display)]
56#[display("an attempt to raise subtree above its depth.")]
57pub struct RaiseAboveRoot;
58
59/// Error indicating that the tree contains just a single known root node and
60/// can't be split into two parts.
61#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Error, Display)]
62#[display("tree contains just a single known root node and can't be split into two parts.")]
63pub struct UnsplittableTree;
64
65/// Error happening when taproot script tree is not complete at certain node.
66#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Error, Display)]
67#[display("taproot script tree is not complete at node {0:?}.")]
68pub struct IncompleteTreeError<N>(N)
69where
70    N: Node + Debug;
71
72/// Errors happening during tree instill operation (see
73/// [`TaprootScriptTree::instill`]).
74#[derive(
75    Clone, PartialOrd, Ord, PartialEq, Eq, Hash, Debug, Display, Error, From
76)]
77#[display(doc_comments)]
78pub enum InstillError {
79    /// unable to instill subtree into taproot script tree since the depth of
80    /// the resulting tree exceeds taproot limit.
81    #[from(MaxDepthExceeded)]
82    MaxDepthExceeded,
83
84    /// unable to instill subtree into taproot script tree since {0}
85    #[from]
86    DfsTraversal(DfsTraversalError),
87}
88
89/// Errors happening during tree cut operation (see [`TaprootScriptTree::cut`]).
90#[derive(
91    Clone, PartialOrd, Ord, PartialEq, Eq, Hash, Debug, Display, Error, From
92)]
93#[display(doc_comments)]
94pub enum CutError {
95    /// unable to instill subtree into taproot script tree since the cut point
96    /// contains leaf or hidden node and thus can't be split into two subtrees.
97    #[from(UnsplittableTree)]
98    UnsplittableTree,
99
100    /// unable to cut subtree from taproot script tree since {0}
101    #[from]
102    DfsTraversal(DfsTraversalError),
103}
104
105/// Error happening when a provided DFS path does not exist within a known part
106/// of a tree.
107#[derive(Clone, PartialOrd, Ord, PartialEq, Eq, Hash, Debug, Display, Error)]
108#[display(doc_comments)]
109pub enum DfsTraversalError {
110    /// the provided DFS path {0} does not exist within a given tree.
111    PathNotExists(DfsPath),
112
113    /// the provided DFS path traverses hidden node {node_hash} at
114    /// {failed_path} to {path_leftover}.
115    HiddenNode {
116        /// The hash of the hidden node found during the path traversal.
117        node_hash: TapNodeHash,
118        /// The path segment which leads to the hidden node.
119        failed_path: DfsPath,
120        /// The path segment which was not able to traverse after the hidden
121        /// node.
122        path_leftover: DfsPath,
123    },
124
125    /// the provided DFS path traverses leaf node {leaf_script} at
126    /// {failed_path} to {path_leftover}.
127    LeafNode {
128        /// The hash of the leaf script of a leaf node found during the path
129        /// traversal.
130        leaf_script: LeafScript,
131        /// The path segment which leads to the leaf node.
132        failed_path: DfsPath,
133        /// The path segment which was not able to traverse after the leaf node.
134        path_leftover: DfsPath,
135    },
136}
137
138/// Represents position of a child node under some parent in DFS (deep first
139/// search) order.
140#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display)]
141#[derive(StrictEncode, StrictDecode)]
142#[strict_encoding(by_order, repr = u8)]
143#[cfg_attr(
144    feature = "serde",
145    derive(Serialize, Deserialize),
146    serde(crate = "serde_crate")
147)]
148pub enum DfsOrder {
149    /// The child node is the first one, in terms of DFS ordering.
150    #[display("dfs-first")]
151    First,
152
153    /// The child node is the last one (i.e. the second one), in terms of DFS
154    /// ordering.
155    #[display("dfs-last")]
156    Last,
157}
158
159impl Not for DfsOrder {
160    type Output = DfsOrder;
161
162    fn not(self) -> Self::Output {
163        match self {
164            DfsOrder::First => DfsOrder::Last,
165            DfsOrder::Last => DfsOrder::First,
166        }
167    }
168}
169
170/// Keeps information about DFS ordering of the child nodes under some parent
171/// node. Used in situations when the node organizes child elements basing on
172/// the lexicographic ordering of the node hashes; but still need to keep
173/// the information about an original DFS ordering.
174#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display)]
175#[derive(StrictEncode, StrictDecode)]
176#[strict_encoding(by_order, repr = u8)]
177#[cfg_attr(
178    feature = "serde",
179    derive(Serialize, Deserialize),
180    serde(crate = "serde_crate")
181)]
182pub enum DfsOrdering {
183    /// The first child under a current ordering is also the first child under
184    /// DFS ordering.
185    #[display("left-to-right")]
186    LeftRight,
187
188    /// The first child under a current ordering is the last child under
189    /// DFS ordering.
190    #[display("right-to-left")]
191    RightLeft,
192}
193
194impl Not for DfsOrdering {
195    type Output = DfsOrdering;
196
197    fn not(self) -> Self::Output {
198        match self {
199            DfsOrdering::LeftRight => DfsOrdering::RightLeft,
200            DfsOrdering::RightLeft => DfsOrdering::LeftRight,
201        }
202    }
203}
204
205/// DFS path within the tree.
206///
207/// A wrapper type around vector of [`DfsOrder`] items for simple display
208/// operations.
209#[derive(
210    Wrapper, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Default, Debug, From
211)]
212#[derive(StrictEncode, StrictDecode)]
213#[cfg_attr(
214    feature = "serde",
215    derive(Serialize, Deserialize),
216    serde(crate = "serde_crate")
217)]
218pub struct DfsPath(Vec<DfsOrder>);
219
220impl AsRef<[DfsOrder]> for DfsPath {
221    #[inline]
222    fn as_ref(&self) -> &[DfsOrder] { self.0.as_ref() }
223}
224
225impl Borrow<[DfsOrder]> for DfsPath {
226    #[inline]
227    fn borrow(&self) -> &[DfsOrder] { self.0.borrow() }
228}
229
230impl Display for DfsPath {
231    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
232        for step in self {
233            f.write_str(match step {
234                DfsOrder::First => "0",
235                DfsOrder::Last => "1",
236            })?;
237        }
238        Ok(())
239    }
240}
241
242/// Error parsing string DFS path representation.
243#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display, Error)]
244#[display("the given DFS path {0} can't be parsed: an unexpected character {1} was found.")]
245pub struct DfsPathParseError(pub String, pub char);
246
247impl FromStr for DfsPath {
248    type Err = DfsPathParseError;
249
250    fn from_str(s: &str) -> Result<Self, Self::Err> {
251        s.chars()
252            .map(|c| match c {
253                '0' => Ok(DfsOrder::First),
254                '1' => Ok(DfsOrder::Last),
255                other => Err(DfsPathParseError(s.to_string(), other)),
256            })
257            .collect()
258    }
259}
260
261impl DfsPath {
262    /// Initializes a new empty path instance.
263    #[inline]
264    pub fn new() -> DfsPath { DfsPath(vec![]) }
265
266    /// Constructs DFS path from an iterator over path steps.
267    pub fn with<'path>(iter: impl IntoIterator<Item = &'path DfsOrder>) -> Self {
268        DfsPath::from_iter(iter)
269    }
270}
271
272impl<'path> IntoIterator for &'path DfsPath {
273    type Item = DfsOrder;
274    type IntoIter = core::iter::Cloned<core::slice::Iter<'path, DfsOrder>>;
275
276    fn into_iter(self) -> Self::IntoIter { self.0.iter().cloned() }
277}
278
279impl IntoIterator for DfsPath {
280    type Item = DfsOrder;
281    type IntoIter = std::vec::IntoIter<DfsOrder>;
282
283    fn into_iter(self) -> Self::IntoIter { self.0.into_iter() }
284}
285
286impl FromIterator<DfsOrder> for DfsPath {
287    fn from_iter<T: IntoIterator<Item = DfsOrder>>(iter: T) -> Self {
288        Self::from_inner(iter.into_iter().collect())
289    }
290}
291
292impl<'iter> FromIterator<&'iter DfsOrder> for DfsPath {
293    fn from_iter<T: IntoIterator<Item = &'iter DfsOrder>>(iter: T) -> Self {
294        Self::from_inner(iter.into_iter().copied().collect())
295    }
296}
297
298/// Trait for taproot tree branch types.
299///
300/// Tree branch is a set of two child nodes.
301pub trait Branch {
302    /// Returns the depth of the subtree under this branch node, if the subtree
303    /// is fully known (i.e. does not contain hidden nodes), or `None`
304    /// otherwise. The depth of subtree for leaf nodes is zero.
305    fn subtree_depth(&self) -> Option<u8>;
306    /// Returns correspondence between internal child node ordering and their
307    /// DFS ordering.
308    fn dfs_ordering(&self) -> DfsOrdering;
309    /// Computes branch hash of this branch node.
310    fn branch_hash(&self) -> TapBranchHash;
311}
312
313/// Trait for taproot tree node types.
314///
315/// Tree node is either a script leaf node, tree branch node or a hidden node.
316pub trait Node {
317    /// Detects if the node is hidden node, represented just by a hash value.
318    /// It can't be known whether hidden node is a leaf node or a branch node.
319    fn is_hidden(&self) -> bool;
320    /// Detects if the node is branch node (i.e. a node with two child nodes).
321    fn is_branch(&self) -> bool;
322    /// Detects if the node is a script leaf node.
323    fn is_leaf(&self) -> bool;
324    /// Computes universal node hash.
325    fn node_hash(&self) -> TapNodeHash;
326    /// Returns the depth of this node within the tree.
327    fn node_depth(&self) -> u8;
328    /// Returns the depth of the subtree under this node, if the subtree is
329    /// fully known (i.e. does not contain hidden nodes), or `None` otherwise.
330    /// The depth of subtree for leaf nodes is zero.
331    fn subtree_depth(&self) -> Option<u8>;
332}
333
334/// Ordered set of two branches under taptree node.
335#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
336#[derive(StrictEncode, StrictDecode)]
337#[cfg_attr(
338    feature = "serde",
339    derive(Serialize, Deserialize),
340    serde(crate = "serde_crate")
341)]
342pub struct BranchNode {
343    /// The left (in bitcoin consensus lexicographic ordering) child node.
344    left: Box<TreeNode>,
345    /// The right (in bitcoin consensus lexicographic ordering) child node.
346    right: Box<TreeNode>,
347    /// The DFS ordering for the branches used in case at least one of the
348    /// child nodes is hidden or both branches has the same subtree depth.
349    /// Ignored otherwise: a direct measurement of subtree depths is used in
350    /// this case instead.
351    dfs_ordering: DfsOrdering,
352}
353
354impl Branch for BranchNode {
355    #[inline]
356    fn subtree_depth(&self) -> Option<u8> {
357        Some(self.left.subtree_depth()?.max(self.right.subtree_depth()?))
358    }
359
360    fn dfs_ordering(&self) -> DfsOrdering { self.dfs_ordering }
361
362    fn branch_hash(&self) -> TapBranchHash {
363        TapBranchHash::from_node_hashes(
364            self.as_left_node().node_hash(),
365            self.as_right_node().node_hash(),
366        )
367    }
368}
369
370impl BranchNode {
371    pub(self) fn with(first: TreeNode, last: TreeNode) -> Self {
372        let hash1 = first.node_hash();
373        let hash2 = last.node_hash();
374        if hash1 < hash2 {
375            BranchNode {
376                left: Box::new(first),
377                right: Box::new(last),
378                dfs_ordering: DfsOrdering::LeftRight,
379            }
380        } else {
381            BranchNode {
382                left: Box::new(last),
383                right: Box::new(first),
384                dfs_ordering: DfsOrdering::RightLeft,
385            }
386        }
387    }
388
389    /// Splits the structure into the left and right nodes, ordered according
390    /// to bitcoin consensus rules (by the lexicographic order of the node
391    /// hash values).
392    #[inline]
393    pub fn split(self) -> (TreeNode, TreeNode) { (*self.left, *self.right) }
394
395    /// Splits the structure into the left and right nodes, ordered according
396    /// to the original DFS order.
397    #[inline]
398    pub fn split_dfs(self) -> (TreeNode, TreeNode) {
399        match self.dfs_ordering {
400            DfsOrdering::LeftRight => (*self.left, *self.right),
401            DfsOrdering::RightLeft => (*self.right, *self.left),
402        }
403    }
404
405    /// Returns reference for to left (in bitcoin consensus lexicographic
406    /// ordering) child node.
407    #[inline]
408    pub fn as_left_node(&self) -> &TreeNode { &self.left }
409
410    /// Returns reference for to right (in bitcoin consensus lexicographic
411    /// ordering) child node.
412    #[inline]
413    pub fn as_right_node(&self) -> &TreeNode { &self.right }
414
415    /// Returns mutable reference to the left (in bitcoin consensus
416    /// lexicographic ordering) child node.
417    #[inline]
418    pub(self) fn as_left_node_mut(&mut self) -> &mut TreeNode { &mut self.left }
419
420    /// Returns reference to the right (in bitcoin consensus lexicographic
421    /// ordering) child node.
422    #[inline]
423    pub(self) fn as_right_node_mut(&mut self) -> &mut TreeNode { &mut self.right }
424
425    /// Returns reference to the child node at specific DFS `direction`.
426    #[inline]
427    pub fn as_dfs_child_node(&self, direction: DfsOrder) -> &TreeNode {
428        match direction {
429            DfsOrder::First => self.as_dfs_first_node(),
430            DfsOrder::Last => self.as_dfs_last_node(),
431        }
432    }
433
434    /// Returns reference to the first (in DFS ordering) child node.
435    #[inline]
436    pub fn as_dfs_first_node(&self) -> &TreeNode {
437        match self.dfs_ordering() {
438            DfsOrdering::LeftRight => self.as_left_node(),
439            DfsOrdering::RightLeft => self.as_right_node(),
440        }
441    }
442
443    /// Returns reference to the last (in DFS ordering) child node.
444    #[inline]
445    pub fn as_dfs_last_node(&self) -> &TreeNode {
446        match self.dfs_ordering() {
447            DfsOrdering::LeftRight => self.as_right_node(),
448            DfsOrdering::RightLeft => self.as_left_node(),
449        }
450    }
451
452    /// Returns mutable reference for the first (in DFS ordering) child node.
453    #[inline]
454    pub(self) fn as_dfs_first_node_mut(&mut self) -> &mut TreeNode {
455        match self.dfs_ordering() {
456            DfsOrdering::LeftRight => self.as_left_node_mut(),
457            DfsOrdering::RightLeft => self.as_right_node_mut(),
458        }
459    }
460
461    /// Returns mutable reference for the last (in DFS ordering) child node.
462    #[inline]
463    pub(self) fn as_dfs_last_node_mut(&mut self) -> &mut TreeNode {
464        match self.dfs_ordering() {
465            DfsOrdering::LeftRight => self.as_right_node_mut(),
466            DfsOrdering::RightLeft => self.as_left_node_mut(),
467        }
468    }
469}
470
471/// Structure representing any complete node inside taproot script tree.
472#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
473#[derive(StrictEncode, StrictDecode)]
474#[strict_encoding(by_order, repr = u8)]
475#[cfg_attr(
476    feature = "serde",
477    derive(Serialize, Deserialize),
478    serde(crate = "serde_crate")
479)]
480pub enum TreeNode {
481    /// Leaf script node. Keeps depth in the second tuple item.
482    Leaf(LeafScript, u8),
483    /// Hidden node, which may be a branch or a leaf node. Keeps depth in the
484    /// second tuple item.
485    Hidden(TapNodeHash, u8),
486    /// Branch node. Keeps depth in the second tuple item.
487    Branch(BranchNode, u8),
488}
489
490impl strict_encoding::StrictEncode for Box<TreeNode> {
491    fn strict_encode<E: Write>(&self, mut e: E) -> Result<usize, strict_encoding::Error> {
492        // This wierd implementation is required because of bug in rust compiler causing
493        // overflow
494        let s = self.as_ref().strict_serialize()?;
495        e.write_all(&s)?;
496        Ok(s.len())
497    }
498}
499
500impl strict_encoding::StrictDecode for Box<TreeNode> {
501    fn strict_decode<D: Read>(d: D) -> Result<Self, strict_encoding::Error> {
502        TreeNode::strict_decode(d).map(Box::new)
503    }
504}
505
506impl TreeNode {
507    /// Constructs leaf tree node.
508    pub fn with_tap_script(script: TapScript, depth: u8) -> TreeNode {
509        TreeNode::Leaf(LeafScript::tapscript(script), depth)
510    }
511
512    /// Constructs branch node without child information. To provide information
513    /// about child nodes use [`PartialBranchNode::push_child`] method.
514    pub fn with_branch(a: TreeNode, b: TreeNode, depth: u8) -> TreeNode {
515        TreeNode::Branch(BranchNode::with(a, b), depth)
516    }
517
518    /// Returns reference to the inner branch node, or `None` for a leaf and
519    /// hidden nodes.
520    pub fn as_branch(&self) -> Option<&BranchNode> {
521        match self {
522            TreeNode::Branch(branch, _) => Some(branch),
523            _ => None,
524        }
525    }
526
527    /// Returns mutable reference to the inner branch node, or `None` for leaf
528    /// and hidden nodes.
529    pub(self) fn as_branch_mut(&mut self) -> Option<&mut BranchNode> {
530        match self {
531            TreeNode::Branch(branch, _) => Some(branch),
532            _ => None,
533        }
534    }
535
536    /// Returns reference to the inner leaf script, or `None` for a branch and
537    /// hidden nodes.
538    pub fn as_leaf_script(&self) -> Option<&LeafScript> {
539        match self {
540            TreeNode::Leaf(leaf_script, _) => Some(leaf_script),
541            _ => None,
542        }
543    }
544
545    /// Traverses tree using the given `path` argument and returns the node
546    /// reference at the tip of the path.
547    ///
548    /// # Errors
549    ///
550    /// Returns [`DfsTraversalError`] if the path can't be traversed.
551    #[inline]
552    pub fn node_at(&self, path: impl AsRef<[DfsOrder]>) -> Result<&TreeNode, DfsTraversalError> {
553        let mut curr = self;
554        let mut past_steps = vec![];
555        let path = path.as_ref();
556        let mut iter = path.iter();
557        for step in iter.by_ref() {
558            past_steps.push(step);
559            let branch = match curr {
560                TreeNode::Branch(branch, _) => branch,
561                TreeNode::Leaf(leaf_script, _) => {
562                    return Err(DfsTraversalError::LeafNode {
563                        leaf_script: leaf_script.clone(),
564                        failed_path: DfsPath::with(past_steps),
565                        path_leftover: iter.collect(),
566                    })
567                }
568                TreeNode::Hidden(hash, _) => {
569                    return Err(DfsTraversalError::HiddenNode {
570                        node_hash: *hash,
571                        failed_path: DfsPath::with(past_steps),
572                        path_leftover: iter.collect(),
573                    })
574                }
575            };
576            curr = match step {
577                DfsOrder::First => branch.as_dfs_first_node(),
578                DfsOrder::Last => branch.as_dfs_last_node(),
579            };
580        }
581        Ok(curr)
582    }
583
584    /// Traverses tree using the given `path` argument and returns the node
585    /// mutable reference at the tip of the path.
586    ///
587    /// # Errors
588    ///
589    /// Returns [`DfsTraversalError`] if the path can't be traversed.
590    #[inline]
591    pub(self) fn node_mut_at<'path>(
592        &mut self,
593        path: impl IntoIterator<Item = &'path DfsOrder>,
594    ) -> Result<&mut TreeNode, DfsTraversalError> {
595        let mut curr = self;
596        let mut past_steps = vec![];
597        let mut iter = path.into_iter();
598        for step in iter.by_ref() {
599            past_steps.push(step);
600            let branch = match curr {
601                TreeNode::Branch(branch, _) => branch,
602                TreeNode::Leaf(leaf_script, _) => {
603                    return Err(DfsTraversalError::LeafNode {
604                        leaf_script: leaf_script.clone(),
605                        failed_path: DfsPath::with(past_steps),
606                        path_leftover: iter.collect(),
607                    })
608                }
609                TreeNode::Hidden(hash, _) => {
610                    return Err(DfsTraversalError::HiddenNode {
611                        node_hash: *hash,
612                        failed_path: DfsPath::with(past_steps),
613                        path_leftover: iter.collect(),
614                    })
615                }
616            };
617            curr = match step {
618                DfsOrder::First => branch.as_dfs_first_node_mut(),
619                DfsOrder::Last => branch.as_dfs_last_node_mut(),
620            };
621        }
622        Ok(curr)
623    }
624
625    /// Returns iterator over all subnodes on a given path.
626    pub(self) fn nodes_on_path<'node, 'path>(
627        &'node self,
628        path: &'path [DfsOrder],
629    ) -> TreePathIter<'node, 'path> {
630        TreePathIter {
631            next_node: Some(self),
632            full_path: path,
633            remaining_path: path.iter(),
634        }
635    }
636
637    /// Returns iterator over all subnodes for this node.
638    pub(self) fn nodes(&self) -> TreeNodeIter { TreeNodeIter::from(self) }
639
640    pub(self) fn nodes_mut(&mut self) -> TreeNodeIterMut { TreeNodeIterMut::from(self) }
641
642    pub(self) fn lower(&mut self, inc: u8) -> Result<u8, MaxDepthExceeded> {
643        let old_depth = self.node_depth();
644        match self {
645            TreeNode::Leaf(_, depth) | TreeNode::Hidden(_, depth) | TreeNode::Branch(_, depth) => {
646                *depth = depth.checked_add(inc).ok_or(MaxDepthExceeded)?;
647            }
648        }
649        Ok(old_depth)
650    }
651
652    pub(self) fn raise(&mut self, dec: u8) -> Result<u8, RaiseAboveRoot> {
653        let old_depth = self.node_depth();
654        match self {
655            TreeNode::Leaf(_, depth) | TreeNode::Hidden(_, depth) | TreeNode::Branch(_, depth) => {
656                *depth = depth.checked_sub(dec).ok_or(RaiseAboveRoot)?;
657            }
658        }
659        Ok(old_depth)
660    }
661
662    /// Checks that the node and all subnodes has correct consensus ordering:
663    /// left-side branch hash is less or equal than right-side branch hash.
664    pub fn check(&self) -> Result<(), TaprootTreeError> {
665        for (node, dfs_path) in self.nodes() {
666            if let Some(branch) = node.as_branch() {
667                let left_hash = branch.left.node_hash();
668                let right_hash = branch.right.node_hash();
669                if left_hash > right_hash {
670                    return Err(TaprootTreeError {
671                        left_hash,
672                        right_hash,
673                        dfs_path,
674                    });
675                }
676            }
677        }
678        Ok(())
679    }
680}
681
682impl Node for TreeNode {
683    fn is_hidden(&self) -> bool { matches!(self, TreeNode::Hidden(..)) }
684
685    fn is_branch(&self) -> bool { matches!(self, TreeNode::Branch(..)) }
686
687    fn is_leaf(&self) -> bool { matches!(self, TreeNode::Leaf(..)) }
688
689    fn node_hash(&self) -> TapNodeHash {
690        match self {
691            TreeNode::Leaf(leaf_script, _) => leaf_script.tap_leaf_hash().into_node_hash(),
692            TreeNode::Hidden(hash, _) => *hash,
693            TreeNode::Branch(branches, _) => branches.branch_hash().into_node_hash(),
694        }
695    }
696
697    fn node_depth(&self) -> u8 {
698        match self {
699            TreeNode::Leaf(_, depth) | TreeNode::Hidden(_, depth) | TreeNode::Branch(_, depth) => {
700                *depth
701            }
702        }
703    }
704
705    fn subtree_depth(&self) -> Option<u8> {
706        match self {
707            TreeNode::Leaf(_, _) => Some(1),
708            TreeNode::Hidden(_, _) => None,
709            TreeNode::Branch(branch, _) => Some(branch.subtree_depth()? + 1),
710        }
711    }
712}
713
714impl TryFrom<PartialTreeNode> for TreeNode {
715    type Error = IncompleteTreeError<PartialTreeNode>;
716
717    fn try_from(partial_node: PartialTreeNode) -> Result<Self, Self::Error> {
718        Ok(match partial_node {
719            PartialTreeNode::Leaf(leaf_script, depth) => TreeNode::Leaf(leaf_script, depth),
720            ref node @ PartialTreeNode::Branch(ref branch, depth) => TreeNode::with_branch(
721                branch
722                    .first
723                    .as_ref()
724                    .ok_or_else(|| IncompleteTreeError(node.clone()))?
725                    .deref()
726                    .clone()
727                    .try_into()?,
728                branch
729                    .second
730                    .as_ref()
731                    .ok_or_else(|| IncompleteTreeError(node.clone()))?
732                    .deref()
733                    .clone()
734                    .try_into()?,
735                depth,
736            ),
737        })
738    }
739}
740
741impl Display for TreeNode {
742    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
743        for (node, path) in self.nodes() {
744            match node {
745                TreeNode::Leaf(leaf_script, depth) => {
746                    writeln!(f, "{} ({}): {}", path, depth, leaf_script)?;
747                }
748                TreeNode::Hidden(hash, depth) => writeln!(f, "{} ({}): {}", path, depth, hash)?,
749                TreeNode::Branch(_, _) => {}
750            }
751        }
752        Ok(())
753    }
754}
755
756/// Structure representing taproot branch node which does not have a complete
757/// information about its childen.
758#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
759pub struct PartialBranchNode {
760    hash: TapBranchHash,
761    first: Option<Box<PartialTreeNode>>,
762    second: Option<Box<PartialTreeNode>>,
763}
764
765impl Branch for PartialBranchNode {
766    fn subtree_depth(&self) -> Option<u8> {
767        Some(
768            self.first
769                .as_ref()?
770                .subtree_depth()?
771                .max(self.second.as_ref()?.subtree_depth()?),
772        )
773    }
774
775    fn dfs_ordering(&self) -> DfsOrdering {
776        match (
777            self.first
778                .as_ref()
779                .map(Box::as_ref)
780                .and_then(PartialTreeNode::subtree_depth),
781            self.second
782                .as_ref()
783                .map(Box::as_ref)
784                .and_then(PartialTreeNode::subtree_depth),
785        ) {
786            (Some(first), Some(second)) => match first.cmp(&second) {
787                // By default we are always ordered in the same way as children were pushed
788                Ordering::Equal => DfsOrdering::LeftRight,
789                Ordering::Less => DfsOrdering::LeftRight,
790                Ordering::Greater => DfsOrdering::RightLeft,
791            },
792            // By default we are always ordered in the same way as children were pushed
793            _ => DfsOrdering::LeftRight,
794        }
795    }
796
797    fn branch_hash(&self) -> TapBranchHash { self.hash }
798}
799
800impl PartialBranchNode {
801    /// Constructs partial branch node without child node information using the
802    /// provided node hash data. If the child nodes are not pushed later, this
803    /// will correspond to a hidden tree node.
804    pub fn with(hash: TapBranchHash) -> Self {
805        PartialBranchNode {
806            hash,
807            first: None,
808            second: None,
809        }
810    }
811
812    /// Adds information about next child node into this branch.
813    ///
814    /// # Returns
815    ///
816    /// Mutable reference to the newly added child node, or `None` if the branch
817    /// was already full (i.e. contained both child nodes).
818    pub fn push_child(&mut self, child: PartialTreeNode) -> Option<&mut PartialTreeNode> {
819        let child = Box::new(child);
820        if let Some(first) = &self.first {
821            if first.node_hash() == child.node_hash() {
822                return self.first.as_deref_mut();
823            }
824        } else {
825            self.first = Some(child);
826            return self.first.as_deref_mut();
827        }
828        if let Some(second) = &self.second {
829            if second.node_hash() == child.node_hash() {
830                self.second.as_deref_mut()
831            } else {
832                None
833            }
834        } else {
835            self.second = Some(child);
836            self.second.as_deref_mut()
837        }
838    }
839
840    /// Returns node hash.
841    #[inline]
842    pub fn node_hash(&self) -> TapNodeHash { TapNodeHash::from_inner(self.hash.into_inner()) }
843}
844
845/// Represents information about taproot script tree when some of the branches
846/// are not complete.
847#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
848pub enum PartialTreeNode {
849    /// Leaf script node. Keeps depth in the second tuple item.
850    Leaf(LeafScript, u8),
851    /// Partial branch node (see [`PartialBranchNode`]). Keeps depth in the
852    /// second tuple item.
853    Branch(PartialBranchNode, u8),
854}
855
856impl PartialTreeNode {
857    /// Constructs leaf node.
858    pub fn with_leaf(leaf_version: LeafVersion, script: Script, depth: u8) -> PartialTreeNode {
859        PartialTreeNode::Leaf(LeafScript::with(leaf_version, script.into()), depth)
860    }
861
862    /// Constructs branch node without child information. To provide information
863    /// about child nodes use [`PartialBranchNode::push_child`] method.
864    pub fn with_branch(hash: TapBranchHash, depth: u8) -> PartialTreeNode {
865        PartialTreeNode::Branch(PartialBranchNode::with(hash), depth)
866    }
867
868    /// Returns reference to the inner branch node, or `None` for the leaf
869    /// nodes.
870    pub fn as_branch(&self) -> Option<&PartialBranchNode> {
871        match self {
872            PartialTreeNode::Leaf(_, _) => None,
873            PartialTreeNode::Branch(branch, _) => Some(branch),
874        }
875    }
876
877    /// Returns mutable reference to the inner branch node, or `None` for the
878    /// leaf nodes.
879    pub fn as_branch_mut(&mut self) -> Option<&mut PartialBranchNode> {
880        match self {
881            PartialTreeNode::Leaf(_, _) => None,
882            PartialTreeNode::Branch(branch, _) => Some(branch),
883        }
884    }
885}
886
887impl Node for PartialTreeNode {
888    #[inline]
889    fn is_hidden(&self) -> bool { false }
890
891    fn is_branch(&self) -> bool { matches!(self, PartialTreeNode::Branch(..)) }
892
893    fn is_leaf(&self) -> bool { matches!(self, PartialTreeNode::Leaf(..)) }
894
895    fn node_hash(&self) -> TapNodeHash {
896        match self {
897            PartialTreeNode::Leaf(leaf_script, _) => leaf_script.tap_leaf_hash().into_node_hash(),
898            PartialTreeNode::Branch(branch, _) => branch.node_hash(),
899        }
900    }
901
902    fn node_depth(&self) -> u8 {
903        match self {
904            PartialTreeNode::Leaf(_, depth) | PartialTreeNode::Branch(_, depth) => *depth,
905        }
906    }
907
908    fn subtree_depth(&self) -> Option<u8> {
909        match self {
910            PartialTreeNode::Leaf(_, _) => Some(0),
911            PartialTreeNode::Branch(branch, _) => branch.subtree_depth(),
912        }
913    }
914}
915
916/// Taproot script tree which keeps internal information in a tree data
917/// structure, which can be modified by adding or removing parts of the tree
918/// (subtrees). See [`Self::join`], [`Self::split`], [`Self::instill`],
919/// [`Self::cut`] operations.
920///
921/// The structure can be build out of (or converted into) [`TapTree`] taproot
922/// tree representation, which doesn't have a modifiable tree structure.
923#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display)]
924#[derive(StrictEncode, StrictDecode)]
925#[cfg_attr(
926    feature = "serde",
927    derive(Serialize, Deserialize),
928    serde(crate = "serde_crate")
929)]
930#[display("{root}")]
931pub struct TaprootScriptTree {
932    root: TreeNode,
933}
934
935impl AsRef<TreeNode> for TaprootScriptTree {
936    #[inline]
937    fn as_ref(&self) -> &TreeNode { &self.root }
938}
939
940impl Borrow<TreeNode> for TaprootScriptTree {
941    #[inline]
942    fn borrow(&self) -> &TreeNode { &self.root }
943}
944
945impl BorrowMut<TreeNode> for TaprootScriptTree {
946    #[inline]
947    fn borrow_mut(&mut self) -> &mut TreeNode { &mut self.root }
948}
949
950impl TaprootScriptTree {
951    /// Constructs new script tree from the root node.
952    ///
953    /// # Errors.
954    ///
955    /// If any of the branches under the root node has non-consensus ordering
956    /// of the child nodes (i.e. by lexicographic order of the node hash
957    /// values).
958    #[inline]
959    pub fn with(root: TreeNode) -> Result<TaprootScriptTree, TaprootTreeError> {
960        root.check()?;
961        Ok(TaprootScriptTree { root })
962    }
963
964    /// Experimental API!
965    ///
966    /// Constructs new script tree from the root node, applying fixes to the
967    /// consensus ordering of the child items, if required.
968    ///
969    /// Tries to fix the underlying subtree structure into consensus-defined
970    /// lexicographic ordering of all branch child nodes.
971    #[stability::unstable(reason = "not sufficiently tested")]
972    #[inline]
973    pub fn with_fixes(root: TreeNode) -> TaprootScriptTree {
974        let mut tree = TaprootScriptTree { root };
975        tree.fix();
976        tree
977    }
978
979    /// Returns iterator over known bitcoin_scripts stored in the tree.
980    ///
981    /// NB: the iterator ignores bitcoin_scripts behind hidden nodes. It
982    /// iterates the bitcoin_scripts in DFS (and not consensus) order.
983    #[inline]
984    pub fn scripts(&self) -> TreeScriptIter { TreeScriptIter::from(self) }
985
986    /// Returns iterator over all known nodes of the tree in DFS order.
987    #[inline]
988    pub fn nodes(&self) -> TreeNodeIter { TreeNodeIter::from(self) }
989
990    /// Returns mutable iterator over all known nodes of the tree in DFS order.
991    #[inline]
992    pub(self) fn nodes_mut(&mut self) -> TreeNodeIterMut { TreeNodeIterMut::from(self) }
993
994    /// Returns iterator over all subnodes on a given path.
995    pub fn nodes_on_path<'node, 'path>(
996        &'node self,
997        path: &'path [DfsOrder],
998    ) -> TreePathIter<'node, 'path> {
999        self.root.nodes_on_path(path)
1000    }
1001
1002    /// Traverses tree using the provided path in DFS order and returns the
1003    /// node reference at the tip of the path.
1004    ///
1005    /// # Errors
1006    ///
1007    /// Returns [`DfsTraversalError`] if the path can't be traversed.
1008    #[inline]
1009    pub fn node_at(&self, path: impl AsRef<[DfsOrder]>) -> Result<&TreeNode, DfsTraversalError> {
1010        self.root.node_at(path)
1011    }
1012
1013    /// Traverses tree using the provided path in DFS order and returns the
1014    /// mutable node reference at the tip of the path.
1015    ///
1016    /// # Errors
1017    ///
1018    /// Returns [`DfsTraversalError`] if the path can't be traversed.
1019    #[inline]
1020    pub(self) fn node_mut_at<'path>(
1021        &mut self,
1022        path: impl IntoIterator<Item = &'path DfsOrder>,
1023    ) -> Result<&mut TreeNode, DfsTraversalError> {
1024        self.root.node_mut_at(path)
1025    }
1026
1027    fn update_ancestors_ordering(&mut self, path: impl Borrow<[DfsOrder]>) {
1028        let path = path.borrow();
1029        for step in (0..path.len()).rev() {
1030            let ancestor = self
1031                .node_mut_at(&path[..step])
1032                .expect("the path must be checked to be valid");
1033            let branch = if let Some(branch) = ancestor.as_branch_mut() {
1034                branch
1035            } else {
1036                return;
1037            };
1038            if branch.left.node_hash() > branch.right.node_hash() {
1039                branch.dfs_ordering = !branch.dfs_ordering;
1040                let old_left = branch.as_left_node().clone();
1041                let old_right = branch.as_right_node().clone();
1042                let left = branch.as_left_node_mut();
1043                *left = old_right;
1044                let right = branch.as_right_node_mut();
1045                *right = old_left;
1046            }
1047        }
1048    }
1049
1050    /// Joins two trees together under a new root.
1051    ///
1052    /// Creates a new tree with the root node containing `self` and `other_tree`
1053    /// as its direct children. The `other_tree` is put into `other_dfs_order`
1054    /// side.
1055    #[inline]
1056    pub fn join(
1057        mut self,
1058        other_tree: TaprootScriptTree,
1059        other_dfs_order: DfsOrder,
1060    ) -> Result<TaprootScriptTree, MaxDepthExceeded> {
1061        self.instill(other_tree, [], other_dfs_order)
1062            .map_err(|_| MaxDepthExceeded)?;
1063        Ok(self)
1064    }
1065
1066    /// Splits the tree into two subtrees. Errors if the tree root is hidden or
1067    /// a script leaf.
1068    ///
1069    /// # Returns
1070    ///
1071    /// Two child nodes under the root of the original tree as a new taproot
1072    /// script trees in the original DFS ordering.
1073    pub fn split(self) -> Result<(TaprootScriptTree, TaprootScriptTree), UnsplittableTree> {
1074        self.cut([], DfsOrder::First).map_err(|_| UnsplittableTree)
1075    }
1076
1077    /// Instills `other_tree` as a subtree under provided `path` by creating a
1078    /// new branch node at the `path` and putting `other_tree` on the `dfs_side`
1079    /// of it.
1080    ///
1081    /// # Error
1082    ///
1083    /// Returns [`InstillError`] when the given path can't be traversed or
1084    /// the resulting tree depth exceeds taproot tree depth limit.
1085    pub fn instill(
1086        &mut self,
1087        mut other_tree: TaprootScriptTree,
1088        path: impl AsRef<[DfsOrder]>,
1089        dfs_order: DfsOrder,
1090    ) -> Result<DfsPath, InstillError> {
1091        let path = path.as_ref();
1092        let depth: u8 = path.len().try_into().map_err(|_| MaxDepthExceeded)?;
1093
1094        let instill_point = self.node_mut_at(path)?;
1095        for n in instill_point.nodes_mut() {
1096            n.lower(1)?;
1097        }
1098        for n in other_tree.nodes_mut() {
1099            n.lower(depth.checked_add(1).ok_or(MaxDepthExceeded)?)?;
1100        }
1101        let instill_root = other_tree.into_root_node();
1102        let branch = if dfs_order == DfsOrder::First {
1103            BranchNode::with(instill_root, instill_point.clone())
1104        } else {
1105            BranchNode::with(instill_point.clone(), instill_root)
1106        };
1107        *instill_point = TreeNode::Branch(branch, depth);
1108
1109        // Update DFS ordering of the nodes above
1110        self.update_ancestors_ordering(path);
1111
1112        let mut path = DfsPath::with(path);
1113        path.push(dfs_order);
1114
1115        Ok(path)
1116    }
1117
1118    /// Cuts subtree out of this tree at the `path`, returning this tree without
1119    /// the cut branch and the cut subtree as a new tree.
1120    ///
1121    /// # Returns
1122    ///
1123    /// Modified original tree without the cut node and a new tree constructed
1124    /// out of the cut node.
1125    ///
1126    /// # Error
1127    ///
1128    /// Returns [`DfsTraversalError`] when the given path can't be traversed or
1129    /// points at an unsplittable node (leaf node or a hidden node).
1130    pub fn cut(
1131        mut self,
1132        path: impl AsRef<[DfsOrder]>,
1133        dfs_side: DfsOrder,
1134    ) -> Result<(TaprootScriptTree, TaprootScriptTree), CutError> {
1135        let path = path.as_ref();
1136        let depth: u8 = path
1137            .len()
1138            .try_into()
1139            .map_err(|_| DfsTraversalError::PathNotExists(path.to_vec().into()))?;
1140
1141        let (mut cut, mut remnant) = match self.node_at(path)? {
1142            TreeNode::Leaf(_, _) | TreeNode::Hidden(_, _) => {
1143                return Err(CutError::UnsplittableTree)
1144            }
1145            TreeNode::Branch(branch, _) if dfs_side == DfsOrder::First => {
1146                branch.clone().split_dfs()
1147            }
1148            TreeNode::Branch(branch, _) => {
1149                let (remnant, cut) = branch.clone().split_dfs();
1150                (cut, remnant)
1151            }
1152        };
1153
1154        for n in cut.nodes_mut() {
1155            n.raise(depth + 1)
1156                .expect("broken taproot tree cut algorithm");
1157        }
1158        for n in remnant.nodes_mut() {
1159            n.raise(1).expect("broken taproot tree cut algorithm");
1160        }
1161
1162        let mut path_iter = path.iter();
1163        if let Some(last_step) = path_iter.next_back() {
1164            let cut_parent = self.node_mut_at(path_iter)?;
1165            let parent_branch_node = cut_parent
1166                .as_branch_mut()
1167                .expect("parent node always a branch node at this point");
1168            let replaced_child = match last_step {
1169                DfsOrder::First => parent_branch_node.as_dfs_first_node_mut(),
1170                DfsOrder::Last => parent_branch_node.as_dfs_last_node_mut(),
1171            };
1172            *replaced_child = remnant;
1173        } else {
1174            self = TaprootScriptTree { root: remnant };
1175        }
1176
1177        let subtree = TaprootScriptTree { root: cut };
1178
1179        // Update DFS ordering of the nodes above
1180        self.update_ancestors_ordering(path);
1181
1182        Ok((self, subtree))
1183    }
1184
1185    /// Returns reference to the root node of the tree.
1186    #[inline]
1187    pub fn as_root_node(&self) -> &TreeNode { &self.root }
1188
1189    /// Consumes the tree and returns instance of the root node of the tree.
1190    #[inline]
1191    pub fn into_root_node(self) -> TreeNode { self.root }
1192
1193    /// Returns a cloned root node.
1194    #[inline]
1195    pub fn to_root_node(&self) -> TreeNode { self.root.clone() }
1196
1197    /// Experimental API!
1198    ///
1199    /// Checks that all nodes in the tree have correct consensus ordering:
1200    /// left-side branch hash is less or equal than right-side branch hash.
1201    #[inline]
1202    #[stability::unstable(
1203        reason = "current stable API assumes that taproot script trees always have correct \
1204                  structure"
1205    )]
1206    pub fn check(&self) -> Result<(), TaprootTreeError> { self.root.check() }
1207
1208    /// Experimental API!
1209    ///
1210    /// Tries to fix the underlying subtree structure into consensus-defined
1211    /// lexicographic ordering of all branch child nodes.
1212    #[stability::unstable(reason = "not sufficiently tested")]
1213    fn fix(&mut self) -> usize {
1214        let mut fix_count = 0usize;
1215        while self.check().is_err() {
1216            let mut path = None;
1217            for (node, p) in self.nodes() {
1218                if node.is_leaf() || node.is_hidden() {
1219                    path = Some(p);
1220                    break;
1221                }
1222            }
1223            if let Some(path) = path {
1224                self.update_ancestors_ordering(path);
1225                fix_count += 1;
1226            }
1227        }
1228        fix_count
1229    }
1230}
1231
1232impl From<TapTree> for TaprootScriptTree {
1233    fn from(tree: TapTree) -> Self {
1234        let mut root: Option<PartialTreeNode> = None;
1235        // TODO: This is a bugfix, which should be reversed once <https://github.com/rust-bitcoin/rust-bitcoin/issues/1069> is fixed upstream
1236        let mut script_leaves = tree.script_leaves().collect::<Vec<_>>();
1237        script_leaves.reverse();
1238        for leaf in script_leaves {
1239            let merkle_branch = leaf.merkle_branch().as_inner();
1240            let leaf_depth = merkle_branch.len() as u8;
1241
1242            let mut curr_hash =
1243                TapLeafHash::from_script(leaf.script(), leaf.leaf_version()).into_node_hash();
1244            let merkle_branch = merkle_branch
1245                .iter()
1246                .map(|step| {
1247                    curr_hash = TapBranchHash::from_node_hashes(*step, curr_hash).into_node_hash();
1248                    curr_hash
1249                })
1250                .collect::<Vec<_>>();
1251            let mut hash_iter = merkle_branch.iter().rev();
1252
1253            match (root.is_some(), hash_iter.next()) {
1254                (false, None) => {
1255                    root = Some(PartialTreeNode::with_leaf(
1256                        leaf.leaf_version(),
1257                        leaf.script().clone(),
1258                        0,
1259                    ))
1260                }
1261                (false, Some(hash)) => {
1262                    root = Some(PartialTreeNode::with_branch(
1263                        TapBranchHash::from_inner(hash.into_inner()),
1264                        0,
1265                    ))
1266                }
1267                (true, None) => unreachable!("broken TapTree structure"),
1268                (true, Some(_)) => {}
1269            }
1270            let mut node = root.as_mut().expect("unreachable");
1271            for (depth, hash) in hash_iter.enumerate() {
1272                match node {
1273                    PartialTreeNode::Leaf(..) => unreachable!("broken TapTree structure"),
1274                    PartialTreeNode::Branch(branch, _) => {
1275                        let child = PartialTreeNode::with_branch(
1276                            TapBranchHash::from_inner(hash.into_inner()),
1277                            depth as u8 + 1,
1278                        );
1279                        node = branch.push_child(child).expect("broken TapTree structure");
1280                    }
1281                }
1282            }
1283            let leaf =
1284                PartialTreeNode::with_leaf(leaf.leaf_version(), leaf.script().clone(), leaf_depth);
1285            match node {
1286                PartialTreeNode::Leaf(..) => { /* nothing to do here */ }
1287                PartialTreeNode::Branch(branch, _) => {
1288                    branch.push_child(leaf);
1289                }
1290            }
1291        }
1292
1293        let root = root
1294            .map(TreeNode::try_from)
1295            .transpose()
1296            .ok()
1297            .flatten()
1298            .expect("broken TapTree structure");
1299
1300        TaprootScriptTree { root }
1301    }
1302}
1303
1304/// Iterator over tree nodes on a path.
1305pub struct TreePathIter<'tree, 'path> {
1306    next_node: Option<&'tree TreeNode>,
1307    full_path: &'path [DfsOrder],
1308    remaining_path: core::slice::Iter<'path, DfsOrder>,
1309}
1310
1311impl<'tree, 'path> Iterator for TreePathIter<'tree, 'path> {
1312    type Item = Result<&'tree TreeNode, DfsTraversalError>;
1313
1314    fn next(&mut self) -> Option<Self::Item> {
1315        match (self.next_node, self.remaining_path.next()) {
1316            (Some(curr_node), Some(step)) => {
1317                match curr_node.node_at([*step]) {
1318                    Err(err) => return Some(Err(err)),
1319                    Ok(next_node) => self.next_node = Some(next_node),
1320                }
1321                Some(Ok(curr_node))
1322            }
1323            (Some(curr_node), None) => {
1324                self.next_node = None;
1325                Some(Ok(curr_node))
1326            }
1327            (None, None) => None,
1328            (None, Some(_)) => Some(Err(DfsTraversalError::PathNotExists(DfsPath::with(
1329                self.full_path,
1330            )))),
1331        }
1332    }
1333}
1334
1335/// Iterator over tree nodes.
1336pub struct TreeNodeIter<'tree> {
1337    stack: Vec<(&'tree TreeNode, DfsPath)>,
1338}
1339
1340impl<'tree, T> From<&'tree T> for TreeNodeIter<'tree>
1341where
1342    T: Borrow<TreeNode>,
1343{
1344    fn from(tree: &'tree T) -> Self {
1345        TreeNodeIter {
1346            stack: vec![(tree.borrow(), DfsPath::new())],
1347        }
1348    }
1349}
1350
1351impl<'tree> Iterator for TreeNodeIter<'tree> {
1352    type Item = (&'tree TreeNode, DfsPath);
1353
1354    fn next(&mut self) -> Option<Self::Item> {
1355        let (curr, path) = self.stack.pop()?;
1356        if let TreeNode::Branch(branch, _) = curr {
1357            let mut p = path.clone();
1358            p.push(DfsOrder::First);
1359            self.stack.push((branch.as_dfs_first_node(), p.clone()));
1360            p.pop();
1361            p.push(DfsOrder::Last);
1362            self.stack.push((branch.as_dfs_last_node(), p));
1363        }
1364        Some((curr, path))
1365    }
1366}
1367
1368struct TreeNodeIterMut<'tree> {
1369    root: &'tree mut TreeNode,
1370    stack: Vec<Vec<DfsOrder>>,
1371}
1372
1373impl<'tree, T> From<&'tree mut T> for TreeNodeIterMut<'tree>
1374where
1375    T: BorrowMut<TreeNode>,
1376{
1377    fn from(tree: &'tree mut T) -> Self {
1378        TreeNodeIterMut {
1379            root: tree.borrow_mut(),
1380            stack: vec![vec![]],
1381        }
1382    }
1383}
1384
1385impl<'tree> Iterator for TreeNodeIterMut<'tree> {
1386    type Item = &'tree mut TreeNode;
1387
1388    fn next(&mut self) -> Option<Self::Item> {
1389        let mut path = self.stack.pop()?;
1390
1391        // We need this because of rust compiler not accepting the fact that
1392        // the root is a part of the self, and that 'tree lifetime will never
1393        // outlive the lifetime of the self.
1394        let mut curr = unsafe { &mut *(self.root as *mut TreeNode) as &'tree mut TreeNode };
1395        for step in &path {
1396            let branch = match curr {
1397                TreeNode::Branch(branch, _) => branch,
1398                _ => unreachable!("iteration algorithm is broken"),
1399            };
1400            curr = match step {
1401                DfsOrder::First => branch.as_dfs_first_node_mut(),
1402                DfsOrder::Last => branch.as_dfs_last_node_mut(),
1403            };
1404        }
1405
1406        if curr.is_branch() {
1407            path.push(DfsOrder::First);
1408            self.stack.push(path.clone());
1409            path.pop();
1410            path.push(DfsOrder::Last);
1411            self.stack.push(path);
1412        }
1413        Some(curr)
1414    }
1415}
1416
1417#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
1418enum BranchDirection {
1419    Shallow,
1420    Deep,
1421}
1422
1423/// Iterator over leaf bitcoin_scripts stored in the leaf nodes of the taproot
1424/// script tree.
1425///
1426/// NB: The bitcoin_scripts are iterated in the DFS order (not consensus).
1427pub struct TreeScriptIter<'tree> {
1428    // Here we store vec of path elements, where each element is a tuple, consisting of:
1429    // 1. Tree node on the path
1430    // 2. Selection of the current branch (false - shallow, true - deep)
1431    path: Vec<(&'tree TreeNode, BranchDirection)>,
1432}
1433
1434impl<'tree, T> From<&'tree T> for TreeScriptIter<'tree>
1435where
1436    T: Borrow<TreeNode>,
1437{
1438    fn from(tree: &'tree T) -> Self {
1439        TreeScriptIter {
1440            path: vec![(tree.borrow(), BranchDirection::Shallow)],
1441        }
1442    }
1443}
1444
1445impl<'tree> Iterator for TreeScriptIter<'tree> {
1446    type Item = (u8, &'tree LeafScript);
1447
1448    fn next(&mut self) -> Option<Self::Item> {
1449        while let Some((node, mut side)) = self.path.pop() {
1450            let mut curr = node;
1451            loop {
1452                match curr {
1453                    // We return only leafs, when found
1454                    TreeNode::Leaf(leaf_script, depth) => {
1455                        return Some((*depth, leaf_script));
1456                    }
1457                    // We skip hidden nodes since we can't do anything about them
1458                    TreeNode::Hidden(..) => break,
1459                    // We restart our search on branching pushing the other
1460                    // branch to the path
1461                    TreeNode::Branch(branch, _) if side == BranchDirection::Shallow => {
1462                        self.path.push((curr, BranchDirection::Deep));
1463                        curr = branch.as_dfs_first_node();
1464                        side = BranchDirection::Shallow;
1465                        continue;
1466                    }
1467                    TreeNode::Branch(branch, _) => {
1468                        curr = branch.as_dfs_last_node();
1469                        side = BranchDirection::Shallow;
1470                        continue;
1471                    }
1472                }
1473            }
1474        }
1475        None
1476    }
1477}
1478
1479impl<'tree> IntoIterator for &'tree TaprootScriptTree {
1480    type Item = (u8, &'tree LeafScript);
1481    type IntoIter = TreeScriptIter<'tree>;
1482
1483    #[inline]
1484    fn into_iter(self) -> Self::IntoIter { self.scripts() }
1485}
1486
1487impl From<&TaprootScriptTree> for TapTree {
1488    fn from(tree: &TaprootScriptTree) -> Self {
1489        let mut builder = TaprootBuilder::new();
1490        for (depth, leaf_script) in tree.scripts() {
1491            builder = builder
1492                .add_leaf_with_ver(depth, leaf_script.script.to_inner(), leaf_script.version)
1493                .expect("broken TaprootScriptTree");
1494        }
1495        TapTree::try_from(builder).expect("broken TaprootScriptTree")
1496    }
1497}
1498
1499impl From<TaprootScriptTree> for TapTree {
1500    #[inline]
1501    fn from(tree: TaprootScriptTree) -> Self { TapTree::from(&tree) }
1502}
1503
1504#[cfg(test)]
1505mod test {
1506    use std::collections::BTreeSet;
1507
1508    use amplify::Wrapper;
1509    use bitcoin::blockdata::opcodes::all;
1510    use bitcoin::hashes::hex::FromHex;
1511    use bitcoin::util::taproot::TaprootBuilder;
1512
1513    use super::*;
1514
1515    /// Composes tree matching a given depth map, filled with dumb script leafs,
1516    /// each of which consists of a single push-int op code, with int value
1517    /// increased for each consecutive leaf.
1518    fn compose_tree(opcode: u8, depth_map: impl IntoIterator<Item = u8>) -> TapTree {
1519        let mut val = opcode;
1520        let mut builder = TaprootBuilder::new();
1521        for depth in depth_map {
1522            let script = Script::from_hex(&format!("{:02x}", val)).unwrap();
1523            builder = builder.add_leaf(depth, script).unwrap();
1524            let (new_val, _) = val.overflowing_add(1);
1525            val = new_val;
1526        }
1527        TapTree::try_from(builder).unwrap()
1528    }
1529
1530    fn test_tree(opcode: u8, depth_map: impl IntoIterator<Item = u8>) {
1531        let taptree = compose_tree(opcode, depth_map);
1532        let script_tree = TaprootScriptTree::from(taptree.clone());
1533
1534        let scripts = taptree
1535            .script_leaves()
1536            .map(|leaf| {
1537                (
1538                    leaf.merkle_branch().as_inner().len() as u8,
1539                    leaf.leaf_version(),
1540                    leaf.script(),
1541                )
1542            })
1543            .collect::<BTreeSet<_>>();
1544        let scripts_prime = script_tree
1545            .scripts()
1546            .map(|(depth, leaf_script)| (depth, leaf_script.version, leaf_script.script.as_inner()))
1547            .collect::<BTreeSet<_>>();
1548        assert_eq!(scripts, scripts_prime);
1549
1550        let taptree_prime = TapTree::from(&script_tree);
1551        assert_eq!(taptree, taptree_prime);
1552    }
1553
1554    fn test_join_split(depth_map: impl IntoIterator<Item = u8>) {
1555        let taptree = compose_tree(0x51, depth_map);
1556        let script_tree = TaprootScriptTree::from(taptree);
1557        assert!(script_tree.check().is_ok());
1558
1559        let instill_tree: TaprootScriptTree = compose_tree(all::OP_RETURN.to_u8(), [0]).into();
1560        let merged_tree = script_tree
1561            .clone()
1562            .join(instill_tree.clone(), DfsOrder::First)
1563            .unwrap();
1564        assert!(merged_tree.check().is_ok());
1565
1566        let _ = TapTree::from(&merged_tree);
1567        assert_ne!(merged_tree, script_tree);
1568
1569        let order = merged_tree.root.as_branch().unwrap().dfs_ordering;
1570
1571        match (
1572            merged_tree.node_at([DfsOrder::First]).unwrap(),
1573            merged_tree.node_at([DfsOrder::Last]).unwrap(),
1574            order,
1575        ) {
1576            (TreeNode::Leaf(leaf_script, 1), _, DfsOrdering::LeftRight)
1577            | (TreeNode::Leaf(leaf_script, 1), _, DfsOrdering::RightLeft)
1578                if leaf_script.script[0] == all::OP_RETURN.to_u8() =>
1579            {
1580                // Everything is fine
1581            }
1582            (_, TreeNode::Leaf(leaf_script, 1), ordering)
1583                if leaf_script.script[0] == all::OP_RETURN.to_u8() =>
1584            {
1585                panic!(
1586                    "instilled tree with script `{:?}` has incorrect DFS ordering {:?}",
1587                    leaf_script.script, ordering
1588                )
1589            }
1590            (TreeNode::Leaf(_, x), _, _) => {
1591                panic!("broken mergged tree depth of first branches: {}", x);
1592            }
1593            _ => panic!("instilled tree is not present as first branch of the merged tree"),
1594        }
1595
1596        let (script_tree_prime, instill_tree_prime) = merged_tree.split().unwrap();
1597        assert!(script_tree_prime.check().is_ok());
1598        assert!(instill_tree_prime.check().is_ok());
1599
1600        assert_eq!(instill_tree, instill_tree_prime);
1601        assert_eq!(script_tree, script_tree_prime);
1602    }
1603
1604    fn test_instill_cut(
1605        depth_map1: impl IntoIterator<Item = u8>,
1606        depth_map2: impl IntoIterator<Item = u8>,
1607        path: &str,
1608    ) {
1609        let path = DfsPath::from_str(path).unwrap();
1610
1611        let taptree = compose_tree(0x51, depth_map1);
1612        let script_tree = TaprootScriptTree::from(taptree);
1613        assert!(script_tree.check().is_ok());
1614
1615        let instill_tree: TaprootScriptTree = compose_tree(50, depth_map2).into();
1616        assert!(instill_tree.check().is_ok());
1617
1618        let mut merged_tree = script_tree.clone();
1619        merged_tree
1620            .instill(instill_tree.clone(), &path, DfsOrder::First)
1621            .unwrap();
1622        assert!(merged_tree.check().is_ok());
1623
1624        let _ = TapTree::from(&merged_tree);
1625        assert_ne!(merged_tree, script_tree);
1626
1627        let (script_tree_prime, instill_tree_prime) =
1628            merged_tree.cut(path, DfsOrder::First).unwrap();
1629
1630        assert!(script_tree_prime.check().is_ok());
1631        assert!(instill_tree_prime.check().is_ok());
1632
1633        assert_eq!(instill_tree, instill_tree_prime);
1634        assert_eq!(script_tree, script_tree_prime);
1635    }
1636
1637    fn testsuite_tree_structures(opcode: u8) {
1638        // Testing all tree variants with up to three levels of depths
1639        // (up to 8 bitcoin_scripts)
1640        test_tree(opcode, [0]);
1641        test_tree(opcode, [1, 1]);
1642        test_tree(opcode, [1, 2, 2]);
1643        test_tree(opcode, [2, 2, 2, 2]);
1644        test_tree(opcode, [1, 2, 3, 3]);
1645        test_tree(opcode, [1, 3, 3, 3, 3]);
1646        // Create a tree as shown below
1647        // A, B , C are at depth 2 and D,E are at 3
1648        //                                       ....
1649        //                                     /      \
1650        //                                    /\      /\
1651        //                                   /  \    /  \
1652        //                                  A    B  C  / \
1653        //                                            D   E
1654        test_tree(opcode, [2, 2, 2, 3, 3]);
1655        test_tree(opcode, [2, 2, 3, 3, 3, 3]);
1656        test_tree(opcode, [2, 3, 3, 3, 3, 3, 3]);
1657        test_tree(opcode, [3, 3, 3, 3, 3, 3, 3, 3]);
1658    }
1659
1660    #[test]
1661    fn taptree_parsing() {
1662        // different opcodes may result in different sorting orders, so we try
1663        // to start with opcodes having different offset
1664        testsuite_tree_structures(0x51);
1665        testsuite_tree_structures(51);
1666        testsuite_tree_structures(0);
1667        testsuite_tree_structures(0x80);
1668    }
1669
1670    #[test]
1671    fn taptree_edge_ops() {
1672        let taptree = compose_tree(0x51, [0]);
1673        let script_tree = TaprootScriptTree::from(taptree);
1674        assert!(script_tree.check().is_ok());
1675        assert_eq!(
1676            script_tree.clone().cut([], DfsOrder::First).unwrap_err(),
1677            CutError::UnsplittableTree
1678        );
1679        assert_eq!(
1680            script_tree.cut([], DfsOrder::Last).unwrap_err(),
1681            CutError::UnsplittableTree
1682        );
1683    }
1684
1685    #[test]
1686    fn taptree_join_split() {
1687        test_join_split([0]);
1688        test_join_split([1, 1]);
1689        test_join_split([1, 2, 2]);
1690        test_join_split([2, 2, 2, 2]);
1691        test_join_split([1, 2, 3, 3]);
1692        test_join_split([1, 3, 3, 3, 3]);
1693        test_join_split([2, 2, 2, 3, 3]);
1694        test_join_split([2, 2, 3, 3, 3, 3]);
1695        test_join_split([2, 3, 3, 3, 3, 3, 3]);
1696        test_join_split([3, 3, 3, 3, 3, 3, 3, 3]);
1697    }
1698
1699    #[test]
1700    fn taptree_instill_cut() {
1701        // Use a tree as shown below for a main tree
1702        // A, B , C are at depth 2 and D, E are at 3
1703        //                                       ....
1704        //                                     /      \
1705        //                                    /\      /\
1706        //                                   /  \    /  \
1707        //                                  A    B  C  / \
1708        //                                            D   E
1709        // Paths to nodes:
1710        // A: 00
1711        // B: 01
1712        // C: 10
1713        // D: 110
1714        // C: 111
1715
1716        // Try instilling a single leaf
1717        test_instill_cut([2, 2, 2, 3, 3], [0], "");
1718        test_instill_cut([2, 2, 2, 3, 3], [0], "0");
1719        test_instill_cut([2, 2, 2, 3, 3], [0], "1");
1720        test_instill_cut([2, 2, 2, 3, 3], [0], "00");
1721        test_instill_cut([2, 2, 2, 3, 3], [0], "01");
1722        test_instill_cut([2, 2, 2, 3, 3], [0], "10");
1723        test_instill_cut([2, 2, 2, 3, 3], [0], "11");
1724        test_instill_cut([2, 2, 2, 3, 3], [0], "110");
1725        test_instill_cut([2, 2, 2, 3, 3], [0], "111");
1726
1727        // Try instilling a subtree
1728        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "");
1729        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "0");
1730        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "1");
1731        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "00");
1732        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "01");
1733        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "10");
1734        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "11");
1735        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "110");
1736        test_instill_cut([2, 2, 2, 3, 3], [1, 2, 3, 3], "111");
1737    }
1738
1739    #[test]
1740    fn instill_path_proof() {
1741        let path = DfsPath::from_str("00101").unwrap();
1742
1743        let taptree = compose_tree(0x51, [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7]);
1744        let script_tree = TaprootScriptTree::from(taptree);
1745        assert!(script_tree.check().is_ok());
1746
1747        let instill_tree: TaprootScriptTree = compose_tree(50, [2, 2, 2, 3, 3]).into();
1748        assert!(instill_tree.check().is_ok());
1749
1750        let mut merged_tree = script_tree;
1751        let instill_path = merged_tree
1752            .instill(instill_tree, &path, DfsOrder::First)
1753            .unwrap();
1754        assert!(merged_tree.check().is_ok());
1755
1756        #[derive(PartialEq, Eq, Debug)]
1757        enum PartnerNode {
1758            Script(String),
1759            Hash(TapNodeHash),
1760        }
1761
1762        let path_partners = merged_tree
1763            .nodes_on_path(&instill_path)
1764            .zip(&instill_path)
1765            .map(|(node, step)| {
1766                let branch = node.unwrap().as_branch().unwrap();
1767                match branch.as_dfs_child_node(!step) {
1768                    TreeNode::Leaf(script, _) => {
1769                        PartnerNode::Script(script.script.as_inner().to_string())
1770                    }
1771                    TreeNode::Hidden(node, _) => PartnerNode::Hash(*node),
1772                    TreeNode::Branch(node, _) => {
1773                        PartnerNode::Hash(node.branch_hash().into_node_hash())
1774                    }
1775                }
1776            })
1777            .collect::<Vec<_>>();
1778
1779        assert_eq!(path_partners, vec![
1780            PartnerNode::Hash(
1781                "e1cc80c5229fa380040f65495b5a7adf102ec6b1bfe51b5c3dbda04ee258529f"
1782                    .parse()
1783                    .unwrap()
1784            ),
1785            PartnerNode::Hash(
1786                "ddad73a07b9a7725185f19d6772b02bd4b3a5525d05afde705c186cdcf588c37"
1787                    .parse()
1788                    .unwrap()
1789            ),
1790            PartnerNode::Script(s!("Script(OP_PUSHNUM_1)")),
1791            PartnerNode::Script(s!("Script(OP_PUSHNUM_4)")),
1792            PartnerNode::Script(s!("Script(OP_PUSHNUM_2)")),
1793            PartnerNode::Script(s!("Script(OP_PUSHNUM_3)")),
1794        ]);
1795    }
1796
1797    #[test]
1798    fn tapscripttree_roudtrip() {
1799        let taptree = compose_tree(0x51, [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7]);
1800        let script_tree = TaprootScriptTree::from(taptree.clone());
1801        let taptree_roundtrip = TapTree::from(script_tree);
1802        assert_eq!(taptree, taptree_roundtrip);
1803    }
1804
1805    #[test]
1806    fn tapscripttree_taptree_eq() {
1807        let taptree = compose_tree(0x51, [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7]);
1808        let script_tree = TaprootScriptTree::from(taptree.clone());
1809        assert!(script_tree.check().is_ok());
1810
1811        // TODO: This is a bugfix, which should be reversed once <https://github.com/rust-bitcoin/rust-bitcoin/issues/1069> is fixed upstream
1812        let mut script_leaves = taptree.script_leaves().collect::<Vec<_>>();
1813        script_leaves.reverse();
1814
1815        for (leaf, (_, leaf_script)) in script_leaves.iter().zip(script_tree.scripts()) {
1816            assert_eq!(leaf.script(), leaf_script.script.as_inner());
1817        }
1818    }
1819
1820    #[test]
1821    fn tapscripttree_dfs() {
1822        let depth_map = [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7];
1823        let mut val = 0x51;
1824
1825        let taptree = compose_tree(val, depth_map);
1826        let script_tree = TaprootScriptTree::from(taptree);
1827        assert!(script_tree.check().is_ok());
1828
1829        for (depth, leaf_script) in script_tree.scripts() {
1830            let script = Script::from_hex(&format!("{:02x}", val)).unwrap();
1831
1832            assert_eq!(depth, depth_map[(val - 0x51) as usize]);
1833            assert_eq!(script, leaf_script.script.to_inner());
1834
1835            let (new_val, _) = val.overflowing_add(1);
1836            val = new_val;
1837        }
1838    }
1839}