From c362d9047321e8d159bb5d1316fae76727206aea Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Thu, 10 Apr 2025 21:29:01 +0000 Subject: [PATCH 1/6] tr: move TapTree depth-checking to construction rather than embedding in Tr This makes it impossible to construct a `TapTree` whose depth exceeds the limit, which means that we can assume that our maximum depth is 128 in all of our code. Also introduces a new error type for this case distinct from "maximum recursive depth exceeded" which is supposed to be about limitations of this library, or at least of Script, rather than about the Taproot limit. --- src/descriptor/mod.rs | 2 +- src/descriptor/tr/mod.rs | 23 +++++++++++------------ src/descriptor/tr/taptree.rs | 24 +++++++++++++++++++++--- src/lib.rs | 9 +++++++++ src/policy/concrete.rs | 6 +++++- src/policy/mod.rs | 17 +++++++++++------ 6 files changed, 58 insertions(+), 23 deletions(-) diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs index 47da23b8b..59a47e890 100644 --- a/src/descriptor/mod.rs +++ b/src/descriptor/mod.rs @@ -42,7 +42,7 @@ pub use self::bare::{Bare, Pkh}; pub use self::segwitv0::{Wpkh, Wsh, WshInner}; pub use self::sh::{Sh, ShInner}; pub use self::sortedmulti::SortedMultiVec; -pub use self::tr::{TapTree, TapTreeIter, TapTreeIterItem, Tr}; +pub use self::tr::{TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem, Tr}; pub mod checksum; mod key; diff --git a/src/descriptor/tr/mod.rs b/src/descriptor/tr/mod.rs index 143376334..22a733ad4 100644 --- a/src/descriptor/tr/mod.rs +++ b/src/descriptor/tr/mod.rs @@ -6,7 +6,7 @@ use core::{cmp, fmt, hash}; use bitcoin::secp256k1; use bitcoin::taproot::{ LeafVersion, TaprootBuilder, TaprootSpendInfo, TAPROOT_CONTROL_BASE_SIZE, - TAPROOT_CONTROL_MAX_NODE_COUNT, TAPROOT_CONTROL_NODE_SIZE, + TAPROOT_CONTROL_NODE_SIZE, }; use bitcoin::{opcodes, Address, Network, ScriptBuf, Weight}; use sync::Arc; @@ -28,7 +28,7 @@ use crate::{ mod taptree; -pub use self::taptree::{TapTree, TapTreeIter, TapTreeIterItem}; +pub use self::taptree::{TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem}; /// A taproot descriptor pub struct Tr { @@ -98,13 +98,7 @@ impl Tr { /// Create a new [`Tr`] descriptor from internal key and [`TapTree`] pub fn new(internal_key: Pk, tree: Option>) -> Result { Tap::check_pk(&internal_key)?; - let nodes = tree.as_ref().map(|t| t.height()).unwrap_or(0); - - if nodes <= TAPROOT_CONTROL_MAX_NODE_COUNT { - Ok(Self { internal_key, tree, spend_info: Mutex::new(None) }) - } else { - Err(Error::MaxRecursiveDepthExceeded) - } + Ok(Self { internal_key, tree, spend_info: Mutex::new(None) }) } /// Obtain the internal key of [`Tr`] descriptor @@ -399,18 +393,23 @@ impl crate::expression::FromTree for Tr { impl<'s, Pk: MiniscriptKey> TreeStack<'s, Pk> { fn new() -> Self { Self { inner: Vec::with_capacity(128) } } - fn push(&mut self, parent: expression::TreeIterItem<'s>, tree: TapTree) { + fn push( + &mut self, + parent: expression::TreeIterItem<'s>, + tree: TapTree, + ) -> Result<(), Error> { let mut next_push = (parent, tree); while let Some(top) = self.inner.pop() { if next_push.0.index() == top.0.index() { next_push.0 = top.0.parent().unwrap(); - next_push.1 = TapTree::combine(top.1, next_push.1); + next_push.1 = TapTree::combine(top.1, next_push.1)?; } else { self.inner.push(top); break; } } self.inner.push(next_push); + Ok(()) } fn pop_final(&mut self) -> Option> { @@ -457,7 +456,7 @@ impl crate::expression::FromTree for Tr { return Err(Error::NonTopLevel(format!("{:?}", script))); }; - tree_stack.push(node.parent().unwrap(), TapTree::Leaf(Arc::new(script))); + tree_stack.push(node.parent().unwrap(), TapTree::Leaf(Arc::new(script)))?; tap_tree_iter.skip_descendants(); } } diff --git a/src/descriptor/tr/taptree.rs b/src/descriptor/tr/taptree.rs index c68033f54..9d73f06f9 100644 --- a/src/descriptor/tr/taptree.rs +++ b/src/descriptor/tr/taptree.rs @@ -2,7 +2,7 @@ use core::{cmp, fmt}; -use bitcoin::taproot::{LeafVersion, TapLeafHash}; +use bitcoin::taproot::{LeafVersion, TapLeafHash, TAPROOT_CONTROL_MAX_NODE_COUNT}; use crate::miniscript::context::Tap; use crate::policy::{Liftable, Semantic}; @@ -10,6 +10,20 @@ use crate::prelude::Vec; use crate::sync::Arc; use crate::{Miniscript, MiniscriptKey, Threshold, ToPublicKey, TranslateErr, Translator}; +/// Tried to construct Taproot tree which was too deep. +#[derive(PartialEq, Eq, Debug)] +#[non_exhaustive] +pub struct TapTreeDepthError; + +impl fmt::Display for TapTreeDepthError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("maximum Taproot tree depth (128) exceeded") + } +} + +#[cfg(feature = "std")] +impl std::error::Error for TapTreeDepthError {} + /// A Taproot Tree representation. // Hidden leaves are not yet supported in descriptor spec. Conceptually, it should // be simple to integrate those here, but it is best to wait on core for the exact syntax. @@ -33,9 +47,13 @@ pub enum TapTree { impl TapTree { /// Creates a `TapTree` by combining `left` and `right` tree nodes. - pub fn combine(left: TapTree, right: TapTree) -> Self { + pub fn combine(left: TapTree, right: TapTree) -> Result { let height = 1 + cmp::max(left.height(), right.height()); - TapTree::Tree { left: Arc::new(left), right: Arc::new(right), height } + if height <= TAPROOT_CONTROL_MAX_NODE_COUNT { + Ok(TapTree::Tree { left: Arc::new(left), right: Arc::new(right), height }) + } else { + Err(TapTreeDepthError) + } } /// Returns the height of this tree. diff --git a/src/lib.rs b/src/lib.rs index 3b0c0d3bf..fe6490d8c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -450,6 +450,8 @@ pub enum Error { LiftError(policy::LiftError), /// Forward script context related errors ContextError(miniscript::context::ScriptContextError), + /// Tried to construct a Taproot tree which was too deep. + TapTreeDepthError(crate::descriptor::TapTreeDepthError), /// Recursion depth exceeded when parsing policy/miniscript from string MaxRecursiveDepthExceeded, /// Anything but c:pk(key) (P2PK), c:pk_h(key) (P2PKH), and thresh_m(k,...) @@ -509,6 +511,7 @@ impl fmt::Display for Error { Error::TypeCheck(ref e) => write!(f, "typecheck: {}", e), Error::Secp(ref e) => fmt::Display::fmt(e, f), Error::ContextError(ref e) => fmt::Display::fmt(e, f), + Error::TapTreeDepthError(ref e) => fmt::Display::fmt(e, f), #[cfg(feature = "compiler")] Error::CompilerError(ref e) => fmt::Display::fmt(e, f), Error::ConcretePolicy(ref e) => fmt::Display::fmt(e, f), @@ -573,6 +576,7 @@ impl std::error::Error for Error { ConcretePolicy(e) => Some(e), LiftError(e) => Some(e), ContextError(e) => Some(e), + TapTreeDepthError(e) => Some(e), AnalysisError(e) => Some(e), PubKeyCtxError(e, _) => Some(e), AbsoluteLockTime(e) => Some(e), @@ -594,6 +598,11 @@ impl From for Error { fn from(e: policy::LiftError) -> Error { Error::LiftError(e) } } +#[doc(hidden)] +impl From for Error { + fn from(e: crate::descriptor::TapTreeDepthError) -> Error { Error::TapTreeDepthError(e) } +} + #[doc(hidden)] impl From for Error { fn from(e: miniscript::context::ScriptContextError) -> Error { Error::ContextError(e) } diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs index 2fe0748b8..36a290e26 100644 --- a/src/policy/concrete.rs +++ b/src/policy/concrete.rs @@ -965,7 +965,11 @@ fn with_huffman_tree(ms: Vec<(OrdF64, Miniscript)>) let (p2, s2) = node_weights.pop().expect("len must at least be two"); let p = (p1.0).0 + (p2.0).0; - node_weights.push((Reverse(OrdF64(p)), TapTree::combine(s1, s2))); + node_weights.push(( + Reverse(OrdF64(p)), + TapTree::combine(s1, s2) + .expect("huffman tree cannot produce depth > 128 given sane weights"), + )); } debug_assert!(node_weights.len() == 1); diff --git a/src/policy/mod.rs b/src/policy/mod.rs index 0b2d0c10a..48f9de45d 100644 --- a/src/policy/mod.rs +++ b/src/policy/mod.rs @@ -393,7 +393,7 @@ mod tests { let left = TapTree::Leaf(left_ms_compilation); let right = TapTree::Leaf(right_ms_compilation); - let tree = TapTree::combine(left, right); + let tree = TapTree::combine(left, right).unwrap(); let expected_descriptor = Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap(); @@ -460,18 +460,23 @@ mod tests { // Arrange leaf compilations (acc. to probabilities) using huffman encoding into a TapTree let tree = TapTree::combine( - TapTree::combine(node_compilations[4].clone(), node_compilations[5].clone()), + TapTree::combine(node_compilations[4].clone(), node_compilations[5].clone()) + .unwrap(), TapTree::combine( TapTree::combine( TapTree::combine( node_compilations[0].clone(), node_compilations[1].clone(), - ), + ) + .unwrap(), node_compilations[3].clone(), - ), + ) + .unwrap(), node_compilations[6].clone(), - ), - ); + ) + .unwrap(), + ) + .unwrap(); let expected_descriptor = Descriptor::new_tr("E".to_string(), Some(tree)).unwrap(); assert_eq!(descriptor, expected_descriptor); From d2e30d5e05afb211baae9a63615ebfad198ccaff Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Thu, 10 Apr 2025 21:46:16 +0000 Subject: [PATCH 2/6] tr: introduce TapTree:leaf constructor This constructor will be part of the new API introduced in the following "flatten TapTree" commit. To reduce the diff of that commit, introduce it now. --- src/descriptor/tr/mod.rs | 2 +- src/descriptor/tr/taptree.rs | 3 +++ src/policy/concrete.rs | 4 ++-- src/policy/mod.rs | 8 ++++---- 4 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/descriptor/tr/mod.rs b/src/descriptor/tr/mod.rs index 22a733ad4..4f391ee53 100644 --- a/src/descriptor/tr/mod.rs +++ b/src/descriptor/tr/mod.rs @@ -456,7 +456,7 @@ impl crate::expression::FromTree for Tr { return Err(Error::NonTopLevel(format!("{:?}", script))); }; - tree_stack.push(node.parent().unwrap(), TapTree::Leaf(Arc::new(script)))?; + tree_stack.push(node.parent().unwrap(), TapTree::leaf(script))?; tap_tree_iter.skip_descendants(); } } diff --git a/src/descriptor/tr/taptree.rs b/src/descriptor/tr/taptree.rs index 9d73f06f9..9c80bbf3f 100644 --- a/src/descriptor/tr/taptree.rs +++ b/src/descriptor/tr/taptree.rs @@ -46,6 +46,9 @@ pub enum TapTree { } impl TapTree { + /// Creates a `TapTree` leaf from a Miniscript. + pub fn leaf>>>(ms: A) -> Self { TapTree::Leaf(ms.into()) } + /// Creates a `TapTree` by combining `left` and `right` tree nodes. pub fn combine(left: TapTree, right: TapTree) -> Result { let height = 1 + cmp::max(left.height(), right.height()); diff --git a/src/policy/concrete.rs b/src/policy/concrete.rs index 36a290e26..ccd666e80 100644 --- a/src/policy/concrete.rs +++ b/src/policy/concrete.rs @@ -589,7 +589,7 @@ impl Policy { .collect() } - /// Gets the number of [TapLeaf](`TapTree::Leaf`)s considering exhaustive root-level [`Policy::Or`] + /// Gets the number of [TapLeaf](`TapTree::leaf`)s considering exhaustive root-level [`Policy::Or`] /// and [`Policy::Thresh`] disjunctions for the `TapTree`. #[cfg(feature = "compiler")] fn num_tap_leaves(&self) -> usize { self.tapleaf_probability_iter().count() } @@ -957,7 +957,7 @@ impl expression::FromTree for Policy { fn with_huffman_tree(ms: Vec<(OrdF64, Miniscript)>) -> TapTree { let mut node_weights = BinaryHeap::<(Reverse, TapTree)>::new(); for (prob, script) in ms { - node_weights.push((Reverse(prob), TapTree::Leaf(Arc::new(script)))); + node_weights.push((Reverse(prob), TapTree::leaf(script))); } assert_ne!(node_weights.len(), 0, "empty Miniscript compilation"); while node_weights.len() > 1 { diff --git a/src/policy/mod.rs b/src/policy/mod.rs index 48f9de45d..55762374e 100644 --- a/src/policy/mod.rs +++ b/src/policy/mod.rs @@ -375,7 +375,7 @@ mod tests { let descriptor = policy.compile_tr(Some(unspendable_key.clone())).unwrap(); let ms_compilation: Miniscript = ms_str!("multi_a(2,A,B,C,D)"); - let tree: TapTree = TapTree::Leaf(Arc::new(ms_compilation)); + let tree: TapTree = TapTree::leaf(ms_compilation); let expected_descriptor = Descriptor::new_tr(unspendable_key.clone(), Some(tree)).unwrap(); assert_eq!(descriptor, expected_descriptor); @@ -391,8 +391,8 @@ mod tests { let right_ms_compilation: Arc> = Arc::new(ms_str!("and_v(v:pk(A),pk(B))")); - let left = TapTree::Leaf(left_ms_compilation); - let right = TapTree::Leaf(right_ms_compilation); + let left = TapTree::leaf(left_ms_compilation); + let right = TapTree::leaf(right_ms_compilation); let tree = TapTree::combine(left, right).unwrap(); let expected_descriptor = @@ -454,7 +454,7 @@ mod tests { .into_iter() .map(|x| { let leaf_policy: Concrete = policy_str!("{}", x); - TapTree::Leaf(Arc::from(leaf_policy.compile::().unwrap())) + TapTree::leaf(leaf_policy.compile::().unwrap()) }) .collect::>(); From 5abc70c6be23d1806d77c509cbf1ed1d9f6077ce Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Thu, 10 Apr 2025 21:54:03 +0000 Subject: [PATCH 3/6] tr: break Tr API in a couple ways We fix the return value of Tr::tap_tree to return a Option<&T> instead of an &Option; we already do this in Descriptor::tap_tree, which is likely to be more commonly called ... and both methods are pretty obscure. Also remove the iter_scripts method in favor of leaves. We did this rename in the previous PR. However, that PR broke its API anyway by making it yield an opaque struct rather than a tuple. Better to just do all the breakage at once. Maaybe we should backport the rename to 12.x as a pure rename, so that people will update their code and get nicer error messages. --- src/descriptor/mod.rs | 6 +++--- src/descriptor/tr/mod.rs | 11 +---------- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs index 59a47e890..8bf75e1d5 100644 --- a/src/descriptor/mod.rs +++ b/src/descriptor/mod.rs @@ -256,7 +256,7 @@ impl Descriptor { /// returned value. pub fn tap_tree(&self) -> Option<&TapTree> { if let Descriptor::Tr(ref tr) = self { - tr.tap_tree().as_ref() + tr.tap_tree() } else { None } @@ -268,7 +268,7 @@ impl Descriptor { /// Taproot descriptor containing only a keyspend, returns an empty iterator. pub fn tap_tree_iter(&self) -> tr::TapTreeIter { if let Descriptor::Tr(ref tr) = self { - if let Some(ref tree) = tr.tap_tree() { + if let Some(tree) = tr.tap_tree() { return tree.leaves(); } } @@ -988,7 +988,7 @@ impl FromStr for Descriptor { // FIXME preserve weird/broken behavior from 12.x. // See https://github.com/rust-bitcoin/rust-miniscript/issues/734 ret.sanity_check()?; - for item in inner.iter_scripts() { + for item in inner.leaves() { item.miniscript() .ext_check(&crate::miniscript::analyzable::ExtParams::sane())?; } diff --git a/src/descriptor/tr/mod.rs b/src/descriptor/tr/mod.rs index 4f391ee53..78ba88dbb 100644 --- a/src/descriptor/tr/mod.rs +++ b/src/descriptor/tr/mod.rs @@ -105,16 +105,7 @@ impl Tr { pub fn internal_key(&self) -> &Pk { &self.internal_key } /// Obtain the [`TapTree`] of the [`Tr`] descriptor - pub fn tap_tree(&self) -> &Option> { &self.tree } - - /// Obtain the [`TapTree`] of the [`Tr`] descriptor - #[deprecated(since = "11.0.0", note = "use tap_tree instead")] - pub fn taptree(&self) -> &Option> { self.tap_tree() } - - /// Iterate over all scripts in merkle tree. If there is no script path, the iterator - /// yields [`None`] - #[deprecated(since = "TBD", note = "use `leaves` instead")] - pub fn iter_scripts(&self) -> TapTreeIter { self.leaves() } + pub fn tap_tree(&self) -> Option<&TapTree> { self.tree.as_ref() } /// Iterates over all the leaves of the tree in depth-first preorder. /// From 552e84408e4ea532a901ff461a0636418758edf9 Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Mon, 14 Apr 2025 19:03:09 +0000 Subject: [PATCH 4/6] tr: flatten TapTree type This replaces the internal representation of `TapTree` with one that matches the BIP371 iterator: it's a vector of depths and leaves and nothing else. It does not track its maximum height (except during construction to return an error if 128 is exceeded); it does not track "internal nodes"; it does not offer an API to iterate through its "internal nodes" except by calling Liftable::lift to get a Semantic copy of the tree. This greatly simplifies the type and many algorithms associated with it. It automatically eliminates recursion from display, lifting, dropping, and so on. As it turns out, we never use TapTrees except to iterate over the leaves, and there is no data about internal nodes that we can even compute, let alone store, so there's no reason to have nodes for them. In later commit/PR we will create a second TapTree-like type, which can be constructed from a TapTree for ToPublicKey keys, and can be used to extract Merkle paths and control blocks. But without ToPublicKey we can't compute these, so TapTree is not the type for that. --- src/descriptor/tr/mod.rs | 9 +- src/descriptor/tr/taptree.rs | 164 ++++++++++++++++------------------- src/policy/mod.rs | 6 +- 3 files changed, 77 insertions(+), 102 deletions(-) diff --git a/src/descriptor/tr/mod.rs b/src/descriptor/tr/mod.rs index 78ba88dbb..660c4878b 100644 --- a/src/descriptor/tr/mod.rs +++ b/src/descriptor/tr/mod.rs @@ -276,7 +276,7 @@ impl Tr { T: Translator, { let tree = match &self.tree { - Some(tree) => Some(tree.translate_helper(translate)?), + Some(tree) => Some(tree.translate_pk(translate)?), None => None, }; let translate_desc = @@ -604,11 +604,4 @@ mod tests { // Note the last ac12 only has ac and fails the predicate assert!(!tr.for_each_key(|k| k.starts_with("acc"))); } - - #[test] - fn height() { - let desc = descriptor(); - let tr = Tr::::from_str(&desc).unwrap(); - assert_eq!(tr.tap_tree().as_ref().unwrap().height(), 2); - } } diff --git a/src/descriptor/tr/taptree.rs b/src/descriptor/tr/taptree.rs index 9c80bbf3f..6ff56e6e7 100644 --- a/src/descriptor/tr/taptree.rs +++ b/src/descriptor/tr/taptree.rs @@ -1,6 +1,6 @@ // SPDX-License-Identifier: CC0-1.0 -use core::{cmp, fmt}; +use core::fmt; use bitcoin::taproot::{LeafVersion, TapLeafHash, TAPROOT_CONTROL_MAX_NODE_COUNT}; @@ -8,7 +8,7 @@ use crate::miniscript::context::Tap; use crate::policy::{Liftable, Semantic}; use crate::prelude::Vec; use crate::sync::Arc; -use crate::{Miniscript, MiniscriptKey, Threshold, ToPublicKey, TranslateErr, Translator}; +use crate::{Miniscript, MiniscriptKey, Threshold, ToPublicKey}; /// Tried to construct Taproot tree which was too deep. #[derive(PartialEq, Eq, Debug)] @@ -25,46 +25,28 @@ impl fmt::Display for TapTreeDepthError { impl std::error::Error for TapTreeDepthError {} /// A Taproot Tree representation. -// Hidden leaves are not yet supported in descriptor spec. Conceptually, it should -// be simple to integrate those here, but it is best to wait on core for the exact syntax. -#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash)] -pub enum TapTree { - /// A taproot tree structure - Tree { - /// Left tree branch. - left: Arc>, - /// Right tree branch. - right: Arc>, - /// Tree height, defined as `1 + max(left_height, right_height)`. - height: usize, - }, - /// A taproot leaf denoting a spending condition - // A new leaf version would require a new Context, therefore there is no point - // in adding a LeafVersion with Leaf type here. All Miniscripts right now - // are of Leafversion::default - Leaf(Arc>), +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TapTree { + depths_leaves: Vec<(u8, Arc>)>, } impl TapTree { /// Creates a `TapTree` leaf from a Miniscript. - pub fn leaf>>>(ms: A) -> Self { TapTree::Leaf(ms.into()) } + pub fn leaf>>>(ms: A) -> Self { + TapTree { depths_leaves: vec![(0, ms.into())] } + } /// Creates a `TapTree` by combining `left` and `right` tree nodes. pub fn combine(left: TapTree, right: TapTree) -> Result { - let height = 1 + cmp::max(left.height(), right.height()); - if height <= TAPROOT_CONTROL_MAX_NODE_COUNT { - Ok(TapTree::Tree { left: Arc::new(left), right: Arc::new(right), height }) - } else { - Err(TapTreeDepthError) - } - } - - /// Returns the height of this tree. - pub fn height(&self) -> usize { - match *self { - TapTree::Tree { left: _, right: _, height } => height, - TapTree::Leaf(..) => 0, + let mut depths_leaves = + Vec::with_capacity(left.depths_leaves.len() + right.depths_leaves.len()); + for (depth, leaf) in left.depths_leaves.iter().chain(right.depths_leaves.iter()) { + if usize::from(*depth) > TAPROOT_CONTROL_MAX_NODE_COUNT - 1 { + return Err(TapTreeDepthError); + } + depths_leaves.push((*depth + 1, Arc::clone(leaf))); } + Ok(Self { depths_leaves }) } /// Iterates over all the leaves of the tree in depth-first preorder. @@ -73,61 +55,73 @@ impl TapTree { /// in the tree, which is the data required by PSBT (BIP 371). pub fn leaves(&self) -> TapTreeIter { TapTreeIter::from_tree(self) } - // Helper function to translate keys - pub(super) fn translate_helper( + /// Converts keys from one type of public key to another. + pub fn translate_pk( &self, - t: &mut T, - ) -> Result, TranslateErr> + translate: &mut T, + ) -> Result, crate::TranslateErr> where - T: Translator, + T: crate::Translator, { - let frag = match *self { - TapTree::Tree { ref left, ref right, ref height } => TapTree::Tree { - left: Arc::new(left.translate_helper(t)?), - right: Arc::new(right.translate_helper(t)?), - height: *height, - }, - TapTree::Leaf(ref ms) => TapTree::Leaf(Arc::new(ms.translate_pk(t)?)), - }; - Ok(frag) + let mut ret = TapTree { depths_leaves: Vec::with_capacity(self.depths_leaves.len()) }; + for (depth, leaf) in &self.depths_leaves { + ret.depths_leaves + .push((*depth, Arc::new(leaf.translate_pk(translate)?))); + } + + Ok(ret) } } impl Liftable for TapTree { fn lift(&self) -> Result, crate::Error> { - fn lift_helper(s: &TapTree) -> Result, crate::Error> { - match *s { - TapTree::Tree { ref left, ref right, height: _ } => Ok(Semantic::Thresh( - Threshold::or(Arc::new(lift_helper(left)?), Arc::new(lift_helper(right)?)), - )), - TapTree::Leaf(ref leaf) => leaf.lift(), - } + let thresh_vec = self + .leaves() + .map(|item| item.miniscript().lift().map(Arc::new)) + .collect::, _>>()?; + let thresh = Threshold::new(1, thresh_vec).expect("no size limit on Semantic threshold"); + Ok(Semantic::Thresh(thresh).normalized()) + } +} + +fn fmt_helper( + view: &TapTree, + f: &mut fmt::Formatter, + mut fmt_ms: impl FnMut(&mut fmt::Formatter, &Miniscript) -> fmt::Result, +) -> fmt::Result { + let mut last_depth = 0; + for item in view.leaves() { + if last_depth > 0 { + f.write_str(",")?; } - let pol = lift_helper(self)?; - Ok(pol.normalized()) + while last_depth < item.depth() { + f.write_str("{")?; + last_depth += 1; + } + fmt_ms(f, item.miniscript())?; + while last_depth > item.depth() { + f.write_str("}")?; + last_depth -= 1; + } + } + + while last_depth > 0 { + f.write_str("}")?; + last_depth -= 1; } + Ok(()) } impl fmt::Display for TapTree { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - TapTree::Tree { ref left, ref right, height: _ } => { - write!(f, "{{{},{}}}", *left, *right) - } - TapTree::Leaf(ref script) => write!(f, "{}", *script), - } + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + fmt_helper(self, f, |f, ms| write!(f, "{}", ms)) } } impl fmt::Debug for TapTree { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - TapTree::Tree { ref left, ref right, height: _ } => { - write!(f, "{{{:?},{:?}}}", *left, *right) - } - TapTree::Leaf(ref script) => write!(f, "{:?}", *script), - } + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + fmt_helper(self, f, |f, ms| write!(f, "{:?}", ms)) } } @@ -145,35 +139,25 @@ impl fmt::Debug for TapTree { /// would yield (2, A), (2, B), (2,C), (3, D), (3, E). /// #[derive(Debug, Clone)] -pub struct TapTreeIter<'a, Pk: MiniscriptKey> { - stack: Vec<(u8, &'a TapTree)>, +pub struct TapTreeIter<'tr, Pk: MiniscriptKey> { + inner: core::slice::Iter<'tr, (u8, Arc>)>, } impl<'tr, Pk: MiniscriptKey> TapTreeIter<'tr, Pk> { /// An empty iterator. - pub fn empty() -> Self { Self { stack: vec![] } } + pub fn empty() -> Self { Self { inner: [].iter() } } /// An iterator over a given tree. - fn from_tree(tree: &'tr TapTree) -> Self { Self { stack: vec![(0, tree)] } } + fn from_tree(tree: &'tr TapTree) -> Self { Self { inner: tree.depths_leaves.iter() } } } -impl<'a, Pk> Iterator for TapTreeIter<'a, Pk> -where - Pk: MiniscriptKey + 'a, -{ - type Item = TapTreeIterItem<'a, Pk>; +impl<'tr, Pk: MiniscriptKey> Iterator for TapTreeIter<'tr, Pk> { + type Item = TapTreeIterItem<'tr, Pk>; fn next(&mut self) -> Option { - while let Some((depth, last)) = self.stack.pop() { - match *last { - TapTree::Tree { ref left, ref right, height: _ } => { - self.stack.push((depth + 1, right)); - self.stack.push((depth + 1, left)); - } - TapTree::Leaf(ref ms) => return Some(TapTreeIterItem { node: ms, depth }), - } - } - None + self.inner + .next() + .map(|&(depth, ref node)| TapTreeIterItem { depth, node }) } } diff --git a/src/policy/mod.rs b/src/policy/mod.rs index 55762374e..64dd380e8 100644 --- a/src/policy/mod.rs +++ b/src/policy/mod.rs @@ -386,10 +386,8 @@ mod tests { let policy: Concrete = policy_str!("or(and(pk(A),pk(B)),and(pk(C),pk(D)))"); let descriptor = policy.compile_tr(Some(unspendable_key.clone())).unwrap(); - let left_ms_compilation: Arc> = - Arc::new(ms_str!("and_v(v:pk(C),pk(D))")); - let right_ms_compilation: Arc> = - Arc::new(ms_str!("and_v(v:pk(A),pk(B))")); + let left_ms_compilation: Miniscript = ms_str!("and_v(v:pk(C),pk(D))"); + let right_ms_compilation: Miniscript = ms_str!("and_v(v:pk(A),pk(B))"); let left = TapTree::leaf(left_ms_compilation); let right = TapTree::leaf(right_ms_compilation); From 960a6cc31b215a9d4def424b4b5c34ffd82af8e8 Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Sat, 29 Mar 2025 00:11:25 +0000 Subject: [PATCH 5/6] tr: optimize parsing of taptrees This gets our benchmarks to roughly where we were before we flattened the data structure. The `_tr_oneleaf_` benchmarks appear to be measurably faster than before but everything else is pretty-much the same. Since the old code was definitely slower -- using Arcs, making tons of tiny allocations, storing internal nodes -- we can infer that actual construction of TapTrees is just not a meaningful portion of the time taken to parse expressions, and further optimization effort should go to the expression parser. --- src/descriptor/tr/mod.rs | 55 +++++++++++++----------------------- src/descriptor/tr/taptree.rs | 55 ++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 35 deletions(-) diff --git a/src/descriptor/tr/mod.rs b/src/descriptor/tr/mod.rs index 660c4878b..c2cadc04a 100644 --- a/src/descriptor/tr/mod.rs +++ b/src/descriptor/tr/mod.rs @@ -377,38 +377,6 @@ impl crate::expression::FromTree for Tr { fn from_tree(root: expression::TreeIterItem) -> Result { use crate::expression::{Parens, ParseTreeError}; - struct TreeStack<'s, Pk: MiniscriptKey> { - inner: Vec<(expression::TreeIterItem<'s>, TapTree)>, - } - - impl<'s, Pk: MiniscriptKey> TreeStack<'s, Pk> { - fn new() -> Self { Self { inner: Vec::with_capacity(128) } } - - fn push( - &mut self, - parent: expression::TreeIterItem<'s>, - tree: TapTree, - ) -> Result<(), Error> { - let mut next_push = (parent, tree); - while let Some(top) = self.inner.pop() { - if next_push.0.index() == top.0.index() { - next_push.0 = top.0.parent().unwrap(); - next_push.1 = TapTree::combine(top.1, next_push.1)?; - } else { - self.inner.push(top); - break; - } - } - self.inner.push(next_push); - Ok(()) - } - - fn pop_final(&mut self) -> Option> { - assert_eq!(self.inner.len(), 1); - self.inner.pop().map(|x| x.1) - } - } - root.verify_toplevel("tr", 1..=2) .map_err(From::from) .map_err(Error::Parse)?; @@ -425,7 +393,7 @@ impl crate::expression::FromTree for Tr { Some(tree) => tree, }; - let mut tree_stack = TreeStack::new(); + let mut tree_builder = taptree::TapTreeBuilder::new(); let mut tap_tree_iter = tap_tree.pre_order_iter(); // while let construction needed because we modify the iterator inside the loop // (by calling skip_descendants to skip over the contents of the tapscripts). @@ -440,6 +408,7 @@ impl crate::expression::FromTree for Tr { node.verify_n_children("taptree branch", 2..=2) .map_err(From::from) .map_err(Error::Parse)?; + tree_builder.push_inner_node()?; } else { let script = Miniscript::from_tree(node)?; // FIXME hack for https://github.com/rust-bitcoin/rust-miniscript/issues/734 @@ -447,11 +416,11 @@ impl crate::expression::FromTree for Tr { return Err(Error::NonTopLevel(format!("{:?}", script))); }; - tree_stack.push(node.parent().unwrap(), TapTree::leaf(script))?; + tree_builder.push_leaf(script); tap_tree_iter.skip_descendants(); } } - Tr::new(internal_key, tree_stack.pop_final()) + Tr::new(internal_key, Some(tree_builder.finalize())) } } @@ -604,4 +573,20 @@ mod tests { // Note the last ac12 only has ac and fails the predicate assert!(!tr.for_each_key(|k| k.starts_with("acc"))); } + + #[test] + fn tr_maximum_depth() { + // Copied from integration tests + let descriptor128 = "tr(X!,{pk(X1!),{pk(X2!),{pk(X3!),{pk(X4!),{pk(X5!),{pk(X6!),{pk(X7!),{pk(X8!),{pk(X9!),{pk(X10!),{pk(X11!),{pk(X12!),{pk(X13!),{pk(X14!),{pk(X15!),{pk(X16!),{pk(X17!),{pk(X18!),{pk(X19!),{pk(X20!),{pk(X21!),{pk(X22!),{pk(X23!),{pk(X24!),{pk(X25!),{pk(X26!),{pk(X27!),{pk(X28!),{pk(X29!),{pk(X30!),{pk(X31!),{pk(X32!),{pk(X33!),{pk(X34!),{pk(X35!),{pk(X36!),{pk(X37!),{pk(X38!),{pk(X39!),{pk(X40!),{pk(X41!),{pk(X42!),{pk(X43!),{pk(X44!),{pk(X45!),{pk(X46!),{pk(X47!),{pk(X48!),{pk(X49!),{pk(X50!),{pk(X51!),{pk(X52!),{pk(X53!),{pk(X54!),{pk(X55!),{pk(X56!),{pk(X57!),{pk(X58!),{pk(X59!),{pk(X60!),{pk(X61!),{pk(X62!),{pk(X63!),{pk(X64!),{pk(X65!),{pk(X66!),{pk(X67!),{pk(X68!),{pk(X69!),{pk(X70!),{pk(X71!),{pk(X72!),{pk(X73!),{pk(X74!),{pk(X75!),{pk(X76!),{pk(X77!),{pk(X78!),{pk(X79!),{pk(X80!),{pk(X81!),{pk(X82!),{pk(X83!),{pk(X84!),{pk(X85!),{pk(X86!),{pk(X87!),{pk(X88!),{pk(X89!),{pk(X90!),{pk(X91!),{pk(X92!),{pk(X93!),{pk(X94!),{pk(X95!),{pk(X96!),{pk(X97!),{pk(X98!),{pk(X99!),{pk(X100!),{pk(X101!),{pk(X102!),{pk(X103!),{pk(X104!),{pk(X105!),{pk(X106!),{pk(X107!),{pk(X108!),{pk(X109!),{pk(X110!),{pk(X111!),{pk(X112!),{pk(X113!),{pk(X114!),{pk(X115!),{pk(X116!),{pk(X117!),{pk(X118!),{pk(X119!),{pk(X120!),{pk(X121!),{pk(X122!),{pk(X123!),{pk(X124!),{pk(X125!),{pk(X126!),{pk(X127!),{pk(X128!),pk(X129)}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}})"; + descriptor128.parse::>().unwrap(); + + // Copied from integration tests + let descriptor129 = "tr(X!,{pk(X1!),{pk(X2!),{pk(X3!),{pk(X4!),{pk(X5!),{pk(X6!),{pk(X7!),{pk(X8!),{pk(X9!),{pk(X10!),{pk(X11!),{pk(X12!),{pk(X13!),{pk(X14!),{pk(X15!),{pk(X16!),{pk(X17!),{pk(X18!),{pk(X19!),{pk(X20!),{pk(X21!),{pk(X22!),{pk(X23!),{pk(X24!),{pk(X25!),{pk(X26!),{pk(X27!),{pk(X28!),{pk(X29!),{pk(X30!),{pk(X31!),{pk(X32!),{pk(X33!),{pk(X34!),{pk(X35!),{pk(X36!),{pk(X37!),{pk(X38!),{pk(X39!),{pk(X40!),{pk(X41!),{pk(X42!),{pk(X43!),{pk(X44!),{pk(X45!),{pk(X46!),{pk(X47!),{pk(X48!),{pk(X49!),{pk(X50!),{pk(X51!),{pk(X52!),{pk(X53!),{pk(X54!),{pk(X55!),{pk(X56!),{pk(X57!),{pk(X58!),{pk(X59!),{pk(X60!),{pk(X61!),{pk(X62!),{pk(X63!),{pk(X64!),{pk(X65!),{pk(X66!),{pk(X67!),{pk(X68!),{pk(X69!),{pk(X70!),{pk(X71!),{pk(X72!),{pk(X73!),{pk(X74!),{pk(X75!),{pk(X76!),{pk(X77!),{pk(X78!),{pk(X79!),{pk(X80!),{pk(X81!),{pk(X82!),{pk(X83!),{pk(X84!),{pk(X85!),{pk(X86!),{pk(X87!),{pk(X88!),{pk(X89!),{pk(X90!),{pk(X91!),{pk(X92!),{pk(X93!),{pk(X94!),{pk(X95!),{pk(X96!),{pk(X97!),{pk(X98!),{pk(X99!),{pk(X100!),{pk(X101!),{pk(X102!),{pk(X103!),{pk(X104!),{pk(X105!),{pk(X106!),{pk(X107!),{pk(X108!),{pk(X109!),{pk(X110!),{pk(X111!),{pk(X112!),{pk(X113!),{pk(X114!),{pk(X115!),{pk(X116!),{pk(X117!),{pk(X118!),{pk(X119!),{pk(X120!),{pk(X121!),{pk(X122!),{pk(X123!),{pk(X124!),{pk(X125!),{pk(X126!),{pk(X127!),{pk(X128!),{pk(X129),pk(X130)}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}})"; + assert!(matches!( + descriptor129 + .parse::>() + .unwrap_err(), + crate::Error::TapTreeDepthError(TapTreeDepthError), + )); + } } diff --git a/src/descriptor/tr/taptree.rs b/src/descriptor/tr/taptree.rs index 6ff56e6e7..8b20deb4f 100644 --- a/src/descriptor/tr/taptree.rs +++ b/src/descriptor/tr/taptree.rs @@ -211,3 +211,58 @@ impl TapTreeIterItem<'_, Pk> { TapLeafHash::from_script(&self.compute_script(), self.leaf_version()) } } + +pub(super) struct TapTreeBuilder { + depths_leaves: Vec<(u8, Arc>)>, + complete_heights: u128, // ArrayVec represented as a bitmap...and a bool. + complete_128: bool, // BIP341 says depths are in [0,128] *inclusive* so 129 possibilities. + current_height: u8, +} + +impl TapTreeBuilder { + pub(super) fn new() -> Self { + Self { + depths_leaves: vec![], + complete_heights: 0, + complete_128: false, + current_height: 0, + } + } + + #[inline] + pub(super) fn push_inner_node(&mut self) -> Result<(), TapTreeDepthError> { + self.current_height += 1; + if usize::from(self.current_height) > TAPROOT_CONTROL_MAX_NODE_COUNT { + return Err(TapTreeDepthError); + } + Ok(()) + } + + #[inline] + pub(super) fn push_leaf>>>(&mut self, ms: A) { + self.depths_leaves.push((self.current_height, ms.into())); + + // Special-case 128 which doesn't fit into the `complete_heights` bitmap + if usize::from(self.current_height) == TAPROOT_CONTROL_MAX_NODE_COUNT { + if self.complete_128 { + self.complete_128 = false; + self.current_height -= 1; + } else { + self.complete_128 = true; + return; + } + } + // Then deal with all other nonzero heights + while self.current_height > 0 { + if self.complete_heights & (1 << self.current_height) == 0 { + self.complete_heights |= 1 << self.current_height; + break; + } + self.complete_heights &= !(1 << self.current_height); + self.current_height -= 1; + } + } + + #[inline] + pub(super) fn finalize(self) -> TapTree { TapTree { depths_leaves: self.depths_leaves } } +} From adf3164bb4eda5ab30236c8f78fcf60df562691d Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Sat, 29 Mar 2025 14:21:42 +0000 Subject: [PATCH 6/6] tr: improve TapTreeIter a bunch TapTreeIter is now a wrapper around core::slice::Iter. So we can propagate several traits: ExactSizeIterator, DoubleEndedIterator and FusedIterator. Annoyingly we cannot propagate TrustedLenIterator since that's nightly-only, nor can we propagate Default which was introduced some time after 1.63.0. core::slice::Iter also implements AsRef<&[T]> which is tempting, but we don't propagate that since it feels like it reveals too much about the internals of TapTreeIter. --- src/descriptor/tr/taptree.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/descriptor/tr/taptree.rs b/src/descriptor/tr/taptree.rs index 8b20deb4f..431c1ccab 100644 --- a/src/descriptor/tr/taptree.rs +++ b/src/descriptor/tr/taptree.rs @@ -161,6 +161,20 @@ impl<'tr, Pk: MiniscriptKey> Iterator for TapTreeIter<'tr, Pk> { } } +impl<'tr, Pk: MiniscriptKey> DoubleEndedIterator for TapTreeIter<'tr, Pk> { + fn next_back(&mut self) -> Option { + self.inner + .next_back() + .map(|&(depth, ref node)| TapTreeIterItem { depth, node }) + } +} + +impl<'tr, Pk: MiniscriptKey> ExactSizeIterator for TapTreeIter<'tr, Pk> { + fn len(&self) -> usize { self.inner.len() } +} + +impl<'tr, Pk: MiniscriptKey> core::iter::FusedIterator for TapTreeIter<'tr, Pk> {} + /// Iterator over all of the leaves of a Taproot tree. /// /// If there is no tree (i.e. this is a keyspend-only Taproot descriptor)