Add methods to count nodes with pre-allocated stack

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -14,6 +14,7 @@ using Reexport
 @reexport import .EquationModule: Node, string_tree, print_tree, copy_node, set_node!
 @reexport import .EquationUtilsModule:
     count_nodes,
+    count_nodes_with_stack,
     count_constants,
     count_depth,
     NodeIndex,

src/EquationUtils.jl

@@ -2,6 +2,44 @@ module EquationUtilsModule
 
 import ..EquationModule: Node, copy_node
 
+"""
+    count_nodes_with_stack(tree::Node{T}, preallocated_stack)::Int where {T}
+
+Count the number of nodes in the tree, using a stack instead of
+recursion. While counting nodes is a quick task already, for further
+speed, using a pre-allocated stack can be signficantly faster,
+especially if you can re-use the same stack for multiple calls.
+
+# Arguments
+- `tree::Node{T}`: The tree to count the nodes of.
+- `preallocated_stack::Vector{Node{T}}`: A pre-allocated stack
+   to use for the counting. This should have a length of the
+   potential max depth of a tree. e.g., you can initialize this
+   with `Array{Node{T}}(undef, 100)` for a max depth of 100.
+"""
+function count_nodes_with_stack(
+    tree::Node{T}, preallocated_stack::Vector{Node{T}}
+)::Int where {T}
+    preallocated_stack[1] = tree
+    count = 0
+    i = 1
+    while i !== 0
+        head = preallocated_stack[i]
+        i -= 1
+        count += 1
+        if head.degree == 1
+            i += 1
+            preallocated_stack[i] = head.l
+        elseif head.degree == 2
+            i += 1
+            preallocated_stack[i] = head.l
+            i += 1
+            preallocated_stack[i] = head.r
+        end
+    end
+    return count
+end
+
 """
     count_nodes(tree::Node{T})::Int where {T}
 

test/test_tree_construction.jl

@@ -38,6 +38,7 @@ for unaop in [cos, exp, safe_log, safe_log2, safe_log10, safe_sqrt, relu, gamma,
 
         # Test Basics
         @test n == 9
+        @test n == count_nodes_with_stack(const_tree, Vector{typeof(const_tree)}(undef, 100))
         @test result == true_result
 
         types_to_test = [Float32, Float64, BigFloat]