Skip to content

Commit 960a6cc

Browse files
apoelstraapoelstra
committed
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.
1 parent 552e844 commit 960a6cc

File tree

2 files changed

+75
-35
lines changed

2 files changed

+75
-35
lines changed

‎src/descriptor/tr/mod.rs

+20-35
Original file line numberDiff line numberDiff line change
@@ -377,38 +377,6 @@ impl<Pk: FromStrKey> crate::expression::FromTree for Tr<Pk> {
377377
fn from_tree(root: expression::TreeIterItem) -> Result<Self, Error> {
378378
use crate::expression::{Parens, ParseTreeError};
379379

380-
struct TreeStack<'s, Pk: MiniscriptKey> {
381-
inner: Vec<(expression::TreeIterItem<'s>, TapTree<Pk>)>,
382-
}
383-
384-
impl<'s, Pk: MiniscriptKey> TreeStack<'s, Pk> {
385-
fn new() -> Self { Self { inner: Vec::with_capacity(128) } }
386-
387-
fn push(
388-
&mut self,
389-
parent: expression::TreeIterItem<'s>,
390-
tree: TapTree<Pk>,
391-
) -> Result<(), Error> {
392-
let mut next_push = (parent, tree);
393-
while let Some(top) = self.inner.pop() {
394-
if next_push.0.index() == top.0.index() {
395-
next_push.0 = top.0.parent().unwrap();
396-
next_push.1 = TapTree::combine(top.1, next_push.1)?;
397-
} else {
398-
self.inner.push(top);
399-
break;
400-
}
401-
}
402-
self.inner.push(next_push);
403-
Ok(())
404-
}
405-
406-
fn pop_final(&mut self) -> Option<TapTree<Pk>> {
407-
assert_eq!(self.inner.len(), 1);
408-
self.inner.pop().map(|x| x.1)
409-
}
410-
}
411-
412380
root.verify_toplevel("tr", 1..=2)
413381
.map_err(From::from)
414382
.map_err(Error::Parse)?;
@@ -425,7 +393,7 @@ impl<Pk: FromStrKey> crate::expression::FromTree for Tr<Pk> {
425393
Some(tree) => tree,
426394
};
427395

428-
let mut tree_stack = TreeStack::new();
396+
let mut tree_builder = taptree::TapTreeBuilder::new();
429397
let mut tap_tree_iter = tap_tree.pre_order_iter();
430398
// while let construction needed because we modify the iterator inside the loop
431399
// (by calling skip_descendants to skip over the contents of the tapscripts).
@@ -440,18 +408,19 @@ impl<Pk: FromStrKey> crate::expression::FromTree for Tr<Pk> {
440408
node.verify_n_children("taptree branch", 2..=2)
441409
.map_err(From::from)
442410
.map_err(Error::Parse)?;
411+
tree_builder.push_inner_node()?;
443412
} else {
444413
let script = Miniscript::from_tree(node)?;
445414
// FIXME hack for https://github.com/rust-bitcoin/rust-miniscript/issues/734
446415
if script.ty.corr.base != crate::miniscript::types::Base::B {
447416
return Err(Error::NonTopLevel(format!("{:?}", script)));
448417
};
449418

450-
tree_stack.push(node.parent().unwrap(), TapTree::leaf(script))?;
419+
tree_builder.push_leaf(script);
451420
tap_tree_iter.skip_descendants();
452421
}
453422
}
454-
Tr::new(internal_key, tree_stack.pop_final())
423+
Tr::new(internal_key, Some(tree_builder.finalize()))
455424
}
456425
}
457426

@@ -604,4 +573,20 @@ mod tests {
604573
// Note the last ac12 only has ac and fails the predicate
605574
assert!(!tr.for_each_key(|k| k.starts_with("acc")));
606575
}
576+
577+
#[test]
578+
fn tr_maximum_depth() {
579+
// Copied from integration tests
580+
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)}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}})";
581+
descriptor128.parse::<crate::Descriptor<String>>().unwrap();
582+
583+
// Copied from integration tests
584+
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)}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}})";
585+
assert!(matches!(
586+
descriptor129
587+
.parse::<crate::Descriptor::<String>>()
588+
.unwrap_err(),
589+
crate::Error::TapTreeDepthError(TapTreeDepthError),
590+
));
591+
}
607592
}

‎src/descriptor/tr/taptree.rs

+55
Original file line numberDiff line numberDiff line change
@@ -211,3 +211,58 @@ impl<Pk: ToPublicKey> TapTreeIterItem<'_, Pk> {
211211
TapLeafHash::from_script(&self.compute_script(), self.leaf_version())
212212
}
213213
}
214+
215+
pub(super) struct TapTreeBuilder<Pk: MiniscriptKey> {
216+
depths_leaves: Vec<(u8, Arc<Miniscript<Pk, Tap>>)>,
217+
complete_heights: u128, // ArrayVec<bool, 129> represented as a bitmap...and a bool.
218+
complete_128: bool, // BIP341 says depths are in [0,128] *inclusive* so 129 possibilities.
219+
current_height: u8,
220+
}
221+
222+
impl<Pk: MiniscriptKey> TapTreeBuilder<Pk> {
223+
pub(super) fn new() -> Self {
224+
Self {
225+
depths_leaves: vec![],
226+
complete_heights: 0,
227+
complete_128: false,
228+
current_height: 0,
229+
}
230+
}
231+
232+
#[inline]
233+
pub(super) fn push_inner_node(&mut self) -> Result<(), TapTreeDepthError> {
234+
self.current_height += 1;
235+
if usize::from(self.current_height) > TAPROOT_CONTROL_MAX_NODE_COUNT {
236+
return Err(TapTreeDepthError);
237+
}
238+
Ok(())
239+
}
240+
241+
#[inline]
242+
pub(super) fn push_leaf<A: Into<Arc<Miniscript<Pk, Tap>>>>(&mut self, ms: A) {
243+
self.depths_leaves.push((self.current_height, ms.into()));
244+
245+
// Special-case 128 which doesn't fit into the `complete_heights` bitmap
246+
if usize::from(self.current_height) == TAPROOT_CONTROL_MAX_NODE_COUNT {
247+
if self.complete_128 {
248+
self.complete_128 = false;
249+
self.current_height -= 1;
250+
} else {
251+
self.complete_128 = true;
252+
return;
253+
}
254+
}
255+
// Then deal with all other nonzero heights
256+
while self.current_height > 0 {
257+
if self.complete_heights & (1 << self.current_height) == 0 {
258+
self.complete_heights |= 1 << self.current_height;
259+
break;
260+
}
261+
self.complete_heights &= !(1 << self.current_height);
262+
self.current_height -= 1;
263+
}
264+
}
265+
266+
#[inline]
267+
pub(super) fn finalize(self) -> TapTree<Pk> { TapTree { depths_leaves: self.depths_leaves } }
268+
}

0 commit comments

Comments
 (0)