Eliminate a bunch more recursion; expand the TreeLike iterator trait a bit

src/iter/mod.rs

@@ -18,7 +18,12 @@ use crate::sync::Arc;
 use crate::{Miniscript, MiniscriptKey, ScriptContext, Terminal};
 
 impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for &'a Miniscript<Pk, Ctx> {
-    fn as_node(&self) -> Tree<Self> {
+    type NaryChildren = &'a [Arc<Miniscript<Pk, Ctx>>];
+
+    fn nary_len(tc: &Self::NaryChildren) -> usize { tc.len() }
+    fn nary_index(tc: Self::NaryChildren, idx: usize) -> Self { Arc::as_ref(&tc[idx]) }
+
+    fn as_node(&self) -> Tree<Self, Self::NaryChildren> {
         use Terminal::*;
         match self.node {
             PkK(..) | PkH(..) | RawPkH(..) | After(..) | Older(..) | Sha256(..) | Hash256(..)
@@ -36,14 +41,19 @@ impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for &'a Miniscript<Pk,
             | OrD(ref left, ref right)
             | OrC(ref left, ref right)
             | OrI(ref left, ref right) => Tree::Binary(left, right),
-            AndOr(ref a, ref b, ref c) => Tree::Nary(Arc::from([a.as_ref(), b, c])),
-            Thresh(ref thresh) => Tree::Nary(thresh.iter().map(Arc::as_ref).collect()),
+            AndOr(ref a, ref b, ref c) => Tree::Ternary(a, b, c),
+            Thresh(ref thresh) => Tree::Nary(thresh.data()),
         }
     }
 }
 
-impl<Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for Arc<Miniscript<Pk, Ctx>> {
-    fn as_node(&self) -> Tree<Self> {
+impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for &'a Arc<Miniscript<Pk, Ctx>> {
+    type NaryChildren = &'a [Arc<Miniscript<Pk, Ctx>>];
+
+    fn nary_len(tc: &Self::NaryChildren) -> usize { tc.len() }
+    fn nary_index(tc: Self::NaryChildren, idx: usize) -> Self { &tc[idx] }
+
+    fn as_node(&self) -> Tree<Self, Self::NaryChildren> {
         use Terminal::*;
         match self.node {
             PkK(..) | PkH(..) | RawPkH(..) | After(..) | Older(..) | Sha256(..) | Hash256(..)
@@ -54,17 +64,45 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for Arc<Miniscript<Pk, Ctx>
             | DupIf(ref sub)
             | Verify(ref sub)
             | NonZero(ref sub)
-            | ZeroNotEqual(ref sub) => Tree::Unary(Arc::clone(sub)),
+            | ZeroNotEqual(ref sub) => Tree::Unary(sub),
+            AndV(ref left, ref right)
+            | AndB(ref left, ref right)
+            | OrB(ref left, ref right)
+            | OrD(ref left, ref right)
+            | OrC(ref left, ref right)
+            | OrI(ref left, ref right) => Tree::Binary(left, right),
+            AndOr(ref a, ref b, ref c) => Tree::Ternary(a, b, c),
+            Thresh(ref thresh) => Tree::Nary(thresh.data()),
+        }
+    }
+}
+
+impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for &'a Terminal<Pk, Ctx> {
+    type NaryChildren = &'a [Arc<Miniscript<Pk, Ctx>>];
+
+    fn nary_len(tc: &Self::NaryChildren) -> usize { tc.len() }
+    fn nary_index(tc: Self::NaryChildren, idx: usize) -> Self { tc[idx].as_inner() }
+
+    fn as_node(&self) -> Tree<Self, Self::NaryChildren> {
+        use Terminal::*;
+        match self {
+            PkK(..) | PkH(..) | RawPkH(..) | After(..) | Older(..) | Sha256(..) | Hash256(..)
+            | Ripemd160(..) | Hash160(..) | True | False | Multi(..) | MultiA(..) => Tree::Nullary,
+            Alt(ref sub)
+            | Swap(ref sub)
+            | Check(ref sub)
+            | DupIf(ref sub)
+            | Verify(ref sub)
+            | NonZero(ref sub)
+            | ZeroNotEqual(ref sub) => Tree::Unary(sub.as_inner()),
             AndV(ref left, ref right)
             | AndB(ref left, ref right)
             | OrB(ref left, ref right)
             | OrD(ref left, ref right)
             | OrC(ref left, ref right)
-            | OrI(ref left, ref right) => Tree::Binary(Arc::clone(left), Arc::clone(right)),
-            AndOr(ref a, ref b, ref c) => {
-                Tree::Nary(Arc::from([Arc::clone(a), Arc::clone(b), Arc::clone(c)]))
-            }
-            Thresh(ref thresh) => Tree::Nary(thresh.iter().map(Arc::clone).collect()),
+            | OrI(ref left, ref right) => Tree::Binary(left.as_inner(), right.as_inner()),
+            AndOr(ref a, ref b, ref c) => Tree::Ternary(a.as_inner(), b.as_inner(), c.as_inner()),
+            Thresh(ref thresh) => Tree::Nary(thresh.data()),
         }
     }
 }

src/iter/tree.rs

@@ -7,21 +7,22 @@
 //!
 
 use crate::prelude::*;
-use crate::sync::Arc;
 
 /// Abstract node of a tree.
 ///
 /// Tracks the arity (out-degree) of a node, which is the only thing that
 /// is needed for iteration purposes.
-pub enum Tree<T> {
+pub enum Tree<T, NT> {
     /// Combinator with no children.
     Nullary,
     /// Combinator with one child.
     Unary(T),
     /// Combinator with two children.
     Binary(T, T),
+    /// Combinator with two children.
+    Ternary(T, T, T),
     /// Combinator with more than two children.
-    Nary(Arc<[T]>),
+    Nary(NT),
 }
 
 /// A trait for any structure which has the shape of a Miniscript tree.
@@ -30,22 +31,36 @@ pub enum Tree<T> {
 /// rather than nodes themselves, because it provides algorithms that
 /// assume copying is cheap.
 ///
-/// To implement this trait, you only need to implement the [`TreeLike::as_node`]
-/// method, which will usually be very mechanical. Everything else is provided.
-/// However, to avoid allocations, it may make sense to also implement
-/// [`TreeLike::n_children`] and [`TreeLike::nth_child`] because the default
-/// implementations will allocate vectors for n-ary nodes.
+/// To implement this trait, you only need to implement the [`TreeLike::as_node`],
+/// [`TreeLike::nary_len`] and `[TreeLike::nary_index'] methods, which should
+/// be very mechanical. Everything else is provided.
 pub trait TreeLike: Clone + Sized {
+    /// An abstraction over the children of n-ary nodes. Typically when
+    /// implementing the trait for `&a T` this will be `&'a [T]`.
+    type NaryChildren: Clone;
+
+    /// Accessor for the length of a [`Self::NaryChildren`].
+    fn nary_len(tc: &Self::NaryChildren) -> usize;
+
+    /// Accessor for a specific child of a [`Self::NaryChildren`].
+    ///
+    /// # Panics
+    ///
+    /// May panic if asked for an element outside of the range
+    /// `0..Self::nary_len(&tc)`.
+    fn nary_index(tc: Self::NaryChildren, idx: usize) -> Self;
+
     /// Interpret the node as an abstract node.
-    fn as_node(&self) -> Tree<Self>;
+    fn as_node(&self) -> Tree<Self, Self::NaryChildren>;
 
     /// Accessor for the number of children this node has.
     fn n_children(&self) -> usize {
         match self.as_node() {
             Tree::Nullary => 0,
             Tree::Unary(..) => 1,
             Tree::Binary(..) => 2,
-            Tree::Nary(children) => children.len(),
+            Tree::Ternary(..) => 3,
+            Tree::Nary(ref children) => Self::nary_len(children),
         }
     }
 
@@ -58,7 +73,14 @@ pub trait TreeLike: Clone + Sized {
             (0, Tree::Binary(sub, _)) => Some(sub),
             (1, Tree::Binary(_, sub)) => Some(sub),
             (_, Tree::Binary(..)) => None,
-            (n, Tree::Nary(children)) => children.get(n).cloned(),
+            (0, Tree::Ternary(sub, _, _)) => Some(sub),
+            (1, Tree::Ternary(_, sub, _)) => Some(sub),
+            (2, Tree::Ternary(_, _, sub)) => Some(sub),
+            (_, Tree::Ternary(..)) => None,
+            (n, Tree::Nary(children)) if n < Self::nary_len(&children) => {
+                Some(Self::nary_index(children, n))
+            }
+            (_, Tree::Nary(..)) => None,
         }
     }
 
@@ -210,8 +232,15 @@ impl<T: TreeLike> Iterator for PreOrderIter<T> {
                 self.stack.push(right);
                 self.stack.push(left);
             }
+            Tree::Ternary(a, b, c) => {
+                self.stack.push(c);
+                self.stack.push(b);
+                self.stack.push(a);
+            }
             Tree::Nary(children) => {
-                self.stack.extend(children.iter().rev().cloned());
+                for i in (0..T::nary_len(&children)).rev() {
+                    self.stack.push(T::nary_index(children.clone(), i));
+                }
             }
         }
         Some(top)