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.

Author: apoelstra

src/descriptor/tr/mod.rs

@@ -377,38 +377,6 @@ impl<Pk: FromStrKey> crate::expression::FromTree for Tr<Pk> {
     fn from_tree(root: expression::TreeIterItem) -> Result<Self, Error> {
         use crate::expression::{Parens, ParseTreeError};
 
-        struct TreeStack<'s, Pk: MiniscriptKey> {
-            inner: Vec<(expression::TreeIterItem<'s>, TapTree<Pk>)>,
-        }
-
-        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<Pk>,
-            ) -> 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<TapTree<Pk>> {
-                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<Pk: FromStrKey> crate::expression::FromTree for Tr<Pk> {
             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,18 +408,19 @@ impl<Pk: FromStrKey> crate::expression::FromTree for Tr<Pk> {
                 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
                 if script.ty.corr.base != crate::miniscript::types::Base::B {
                     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()))
     }
 }
 

src/descriptor/tr/taptree.rs

@@ -211,3 +211,41 @@ impl<Pk: ToPublicKey> TapTreeIterItem<'_, Pk> {
         TapLeafHash::from_script(&self.compute_script(), self.leaf_version())
     }
 }
+
+pub(super) struct TapTreeBuilder<Pk: MiniscriptKey> {
+    depths_leaves: Vec<(u8, Arc<Miniscript<Pk, Tap>>)>,
+    complete_heights: u128, // ArrayVec<bool, 128> represented as a bitmap...
+    current_height: u8,     // ...and a height
+}
+
+impl<Pk: MiniscriptKey> TapTreeBuilder<Pk> {
+    pub(super) fn new() -> Self {
+        Self { depths_leaves: vec![], complete_heights: 0, 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<A: Into<Arc<Miniscript<Pk, Tap>>>>(&mut self, ms: A) {
+        self.depths_leaves.push((self.current_height, ms.into()));
+
+        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<Pk> { TapTree { depths_leaves: self.depths_leaves } }
+}