ArenaNode: traversal-based postfix emit + cursor-based rewrite

Author: MilesCranmerBot

src/ArenaNode.jl

@@ -332,20 +332,62 @@ struct PostfixExpr{T,D}
     op::Vector{UInt8}
 end
 
-"""Emit a [`PostfixExpr`](@ref) for an arena-backed tree.
+"""Emit a [`PostfixExpr`](@ref) (postorder / postfix) encoding of `tree`.
 
-This copies the arena prefix `1:tree.idx`.
+This traverses the tree via child pointers and emits nodes in postorder, with the
+root last.
+
+!!! note
+    This is intended as a **serialization/debug utility**. It is *not* an execution
+    strategy (i.e. we should not repeatedly convert between representations to make
+    algorithms work).
 """
 function emit_postfix(tree::ArenaNode{T,D}) where {T,D}
-    n = Int(tree.idx)
     a = tree.arena
-    return PostfixExpr{T,D}(
-        copy(@view a.degree[1:n]),
-        copy(@view a.constant[1:n]),
-        copy(@view a.val[1:n]),
-        copy(@view a.feature[1:n]),
-        copy(@view a.op[1:n]),
-    )
+
+    # Preallocate using the reachable node count (ignoring potential sharing).
+    n = length(tree; break_sharing=Val(true))
+    degree = UInt8[]
+    constant = Bool[]
+    val = T[]
+    feature = UInt16[]
+    op = UInt8[]
+    sizehint!(degree, n)
+    sizehint!(constant, n)
+    sizehint!(val, n)
+    sizehint!(feature, n)
+    sizehint!(op, n)
+
+    # Iterative postorder traversal: (idx, expanded_children?).
+    stack = Tuple{Int32,Bool}[]
+    sizehint!(stack, n)
+    push!(stack, (tree.idx, false))
+
+    while !isempty(stack)
+        idx, expanded = pop!(stack)
+        i = Int(idx)
+        if expanded
+            @inbounds begin
+                push!(degree, a.degree[i])
+                push!(constant, a.constant[i])
+                push!(val, a.val[i])
+                push!(feature, a.feature[i])
+                push!(op, a.op[i])
+            end
+        else
+            push!(stack, (idx, true))
+            d = @inbounds a.degree[i]
+            if d != 0
+                child_idxs = @inbounds a.children[i]
+                @inbounds for j in Int(d):-1:1
+                    c = child_idxs[j]
+                    c != 0 && push!(stack, (c, false))
+                end
+            end
+        end
+    end
+
+    return PostfixExpr{T,D}(degree, constant, val, feature, op)
 end
 
 @generated function _postfix_pop_children(
@@ -436,13 +478,23 @@ is_commutative(_) = false
 
 """Rewrite commutative binary operators so that constants appear on the right.
 
-This is a minimal in-place rewrite that *preserves postfix validity* since it does
-not change any node degrees; it only swaps child pointers.
+This is a minimal in-place rewrite that only swaps child pointers.
+
+!!! note
+    This does **not** rely on any arena index ordering (e.g. postfix layout). It traverses
+    the tree via child pointers.
+
+Since it does not change degrees, any postfix encoding that depends only on `degree`
+(e.g. [`emit_postfix`](@ref)) is preserved.
 """
 function rewrite_commutative_constants_right!(tree::ArenaNode{T,D}, operators) where {T,D}
-    root_idx = Int(tree.idx)
-    @inbounds for i in 1:root_idx
-        node = ArenaNode{T,D}(tree.arena, i)
+    cursor = ArenaCursor(tree)
+    reset!(cursor, tree)
+
+    while true
+        node = next!(cursor)
+        node === nothing && break
+
         node.degree == 2 || continue
         f = operators.binops[node.op]
         is_commutative(f) || continue
@@ -454,6 +506,7 @@ function rewrite_commutative_constants_right!(tree::ArenaNode{T,D}, operators) w
             set_child!(node, l, 2)
         end
     end
+
     return tree
 end
 

test/test_arenanode.jl

@@ -37,26 +37,6 @@ const AN = DynamicExpressions.ArenaNodeModule
     collected_idxs = map(n -> n.idx, collected)
     @test collected_idxs == seen
 
-    # Postfix stack-based utilities (mirroring symbolic_regression.rs patterns):
-    @test AN.is_valid_postfix(atree)
-
-    sizes = Int[]
-    size_stack = Int[]
-    AN.subtree_sizes_into!(atree, sizes, size_stack)
-    start, stop = AN.subtree_range(sizes, Int(atree.idx))
-    @test start == 1
-    @test stop == Int(atree.idx)
-
-    depth_stack = Int[]
-    depth_postfix = AN.tree_mapreduce_postfix_with_stack(
-        atree,
-        _ -> 1,
-        _ -> 0,
-        (_, children) -> maximum(children) + 1,
-        depth_stack,
-    )
-    @test depth_postfix == count_depth(atree)
-
     # Evaluation should match:
     X = randn(Float64, 1, 50)
     y_tree, ok_tree = eval_tree_array(tree, X, operators)
@@ -65,31 +45,6 @@ const AN = DynamicExpressions.ArenaNodeModule
     @test ok_atree
     @test y_tree ≈ y_atree
 
-    # Postfix roundtrip sanity check (debug utility; not an execution strategy):
-    pf = AN.emit_postfix(atree)
-    atree_pf = AN.parse_postfix_to_arena(pf)
-    @test AN.is_valid_postfix(atree_pf)
-    @test count_nodes(atree_pf) == count_nodes(atree)
-    @test string_tree(atree_pf, operators) == string_tree(atree, operators)
-    y_pf, ok_pf = eval_tree_array(atree_pf, X, operators)
-    @test ok_pf
-    @test y_pf ≈ y_tree
-
-    # Minimal rewrite prototype should preserve postfix validity:
-    tree_constleft = 3.2 * x1
-    atree_constleft = AN.arena_from_tree(tree_constleft)
-    @test AN.is_valid_postfix(atree_constleft)
-    y_before, ok_before = eval_tree_array(atree_constleft, X, operators)
-    @test ok_before
-    @test atree_constleft.l.constant
-    AN.rewrite_commutative_constants_right!(atree_constleft, operators)
-    @test AN.is_valid_postfix(atree_constleft)
-    @test !atree_constleft.l.constant
-    @test atree_constleft.r.constant
-    y_after, ok_after = eval_tree_array(atree_constleft, X, operators)
-    @test ok_after
-    @test y_after ≈ y_before
-
     # In-place set_node! should work even when the source tree is from a different arena.
     # (This is important for API-compat with algorithms that construct new subtrees.)
     atree_setnode = AN.arena_from_tree(tree)
@@ -132,4 +87,51 @@ const AN = DynamicExpressions.ArenaNodeModule
     y_tree2, ok_tree2 = eval_tree_array(tree2, X, operators)
     @test ok_tree2
     @test y_tree2 ≈ y_mut
+
+    @testset "Postfix / debug utilities (not an execution strategy)" begin
+        # Postfix stack-based utilities (mirroring symbolic_regression.rs patterns):
+        @test AN.is_valid_postfix(atree)
+
+        sizes = Int[]
+        size_stack = Int[]
+        AN.subtree_sizes_into!(atree, sizes, size_stack)
+        start, stop = AN.subtree_range(sizes, Int(atree.idx))
+        @test start == 1
+        @test stop == Int(atree.idx)
+
+        depth_stack = Int[]
+        depth_postfix = AN.tree_mapreduce_postfix_with_stack(
+            atree,
+            _ -> 1,
+            _ -> 0,
+            (_, children) -> maximum(children) + 1,
+            depth_stack,
+        )
+        @test depth_postfix == count_depth(atree)
+
+        # Postfix roundtrip sanity check (debug utility; not an execution strategy):
+        pf = AN.emit_postfix(atree)
+        atree_pf = AN.parse_postfix_to_arena(pf)
+        @test AN.is_valid_postfix(atree_pf)
+        @test count_nodes(atree_pf) == count_nodes(atree)
+        @test string_tree(atree_pf, operators) == string_tree(atree, operators)
+        y_pf, ok_pf = eval_tree_array(atree_pf, X, operators)
+        @test ok_pf
+        @test y_pf ≈ y_tree
+
+        # Minimal rewrite prototype should preserve postfix validity:
+        tree_constleft = 3.2 * x1
+        atree_constleft = AN.arena_from_tree(tree_constleft)
+        @test AN.is_valid_postfix(atree_constleft)
+        y_before, ok_before = eval_tree_array(atree_constleft, X, operators)
+        @test ok_before
+        @test atree_constleft.l.constant
+        AN.rewrite_commutative_constants_right!(atree_constleft, operators)
+        @test AN.is_valid_postfix(atree_constleft)
+        @test !atree_constleft.l.constant
+        @test atree_constleft.r.constant
+        y_after, ok_after = eval_tree_array(atree_constleft, X, operators)
+        @test ok_after
+        @test y_after ≈ y_before
+    end
 end