fix: stabilize ArenaNode prototype traversal and tests

Author: MilesCranmerBot

src/ArenaNode.jl

@@ -10,7 +10,9 @@ import ..NodeModule:
     unsafe_get_children,
     get_child,
     set_child!,
-    set_children!
+    set_children!,
+    branch_equal,
+    leaf_equal
 
 """Array-backed arena storing the fields of a tree node in a struct-of-arrays form.
 
@@ -55,10 +57,14 @@ Core fields are accessed and mutated via `getproperty`/`setproperty!`.
 struct ArenaNode{T,D} <: AbstractExpressionNode{T,D}
     arena::Arena{T,D}
     idx::Int32
+
+    @inline function ArenaNode{T,D}(arena::Arena{T,D}, idx::Int32) where {T,D}
+        return new{T,D}(arena, idx)
+    end
 end
 
-@inline ArenaNode(arena::Arena{T,D}, idx::Integer) where {T,D} =
-    ArenaNode{T,D}(arena, Int32(idx))
+@inline ArenaNode(arena::Arena{T,D}, idx::Int32) where {T,D} =
+    ArenaNode{T,D}(arena, idx)
 
 @inline function _zero_children(::Val{D}) where {D}
     return ntuple(_ -> Int32(0), Val(D))
@@ -179,6 +185,37 @@ end
     end
 end
 
+# DispatchDoctor's `@stable` checking is based on type inference from argument *types*,
+# and does not consider constant propagation of `getproperty(x, ::Symbol)`. Define
+# ArenaNode-specific equality helpers that access the underlying arena directly.
+@inline function branch_equal(a::ArenaNode{T,D}, b::ArenaNode{T,D})::Bool where {T,D}
+    arena_a = getfield(a, :arena)
+    arena_b = getfield(b, :arena)
+    ia = Int(getfield(a, :idx))
+    ib = Int(getfield(b, :idx))
+    @inbounds return arena_a.op[ia] == arena_b.op[ib]
+end
+
+@inline function leaf_equal(
+    a::ArenaNode{T1,D}, b::ArenaNode{T2,D}
+)::Bool where {T1,T2,D}
+    arena_a = getfield(a, :arena)
+    arena_b = getfield(b, :arena)
+    ia = Int(getfield(a, :idx))
+    ib = Int(getfield(b, :idx))
+
+    @inbounds begin
+        const_a = arena_a.constant[ia]
+        const_b = arena_b.constant[ib]
+        const_a == const_b || return false
+        if const_a
+            return arena_a.val[ia]::T1 == arena_b.val[ib]::T2
+        else
+            return arena_a.feature[ia] == arena_b.feature[ib]
+        end
+    end
+end
+
 """Return an `NTuple{D,Nullable{ArenaNode}}` of children wrappers.
 
 Unused slots are represented as poison nodes (mirroring `Node`), so that
@@ -225,7 +262,7 @@ end
 end
 
 @inline function set_children!(
-    n::ArenaNode{T,D}, children::Union{Tuple,AbstractVector}
+    n::ArenaNode{T,D}, children::Union{Tuple,AbstractVector{<:AbstractNode{D}}}
 ) where {T,D}
     D2 = length(children)
     idxs = _zero_children(Val(D))
@@ -291,22 +328,22 @@ function arena_from_tree(tree::AbstractExpressionNode{T,D}) where {T,D}
 end
 
 """Convert an arena-backed node back into a heap-allocated `Node` tree."""
-function tree_from_arena(tree::ArenaNode{T,D}) where {T,D}
-    function rebuild(n::ArenaNode{T,D})
-        d = n.degree
-        if d == 0
-            return n.constant ? Node{T,D}(; val=n.val) : Node{T,D}(T; feature=n.feature)
-        else
-            # Use a vector here to avoid `Val(d)` with runtime `d`.
-            cs = Vector{Node{T,D}}(undef, Int(d))
-            @inbounds for i in 1:Int(d)
-                cs[i] = rebuild(get_child(n, i))
-            end
-            return Node{T,D}(T; op=n.op, children=cs)
+@inline function _tree_from_arena(n::ArenaNode{T,D})::Node{T,D} where {T,D}
+    d = n.degree
+    if d == 0
+        return n.constant ? Node{T,D}(; val=n.val) : Node{T,D}(T; feature=n.feature)
+    else
+        # Use a vector here to avoid `Val(d)` with runtime `d`.
+        cs = Vector{Node{T,D}}(undef, Int(d))
+        @inbounds for i in 1:Int(d)
+            cs[i] = _tree_from_arena(get_child(n, i))
         end
+        return Node{T,D}(T; op=n.op, children=cs)
     end
+end
 
-    return rebuild(tree)
+function tree_from_arena(tree::ArenaNode{T,D}) where {T,D}
+    return _tree_from_arena(tree)
 end
 
 ################################################################################
@@ -492,8 +529,9 @@ function rewrite_commutative_constants_right!(tree::ArenaNode{T,D}, operators) w
     reset!(cursor, tree)
 
     while true
-        node = next!(cursor)
-        node === nothing && break
+        maybe_node = next!(cursor)
+        maybe_node.null && break
+        node = maybe_node[]
 
         node.degree == 2 || continue
         f = operators.binops[node.op]
@@ -637,7 +675,7 @@ function tree_mapreduce_postfix_with_stack(
     sizehint!(stack, root_idx)
 
     @inbounds for i in 1:root_idx
-        node = ArenaNode(tree.arena, i)
+        node = ArenaNode(tree.arena, Int32(i))
         d = node.degree
         if d == 0
             push!(stack, f_leaf(node)::R)
@@ -679,8 +717,9 @@ mutable struct ArenaCursor{T,D}
     end
 end
 
-@inline ArenaCursor(tree::ArenaNode{T,D}; capacity::Integer=0) where {T,D} =
-    ArenaCursor(tree.arena; capacity)
+@inline function ArenaCursor(tree::ArenaNode{T,D}; capacity::Integer=0) where {T,D}
+    return ArenaCursor(tree.arena; capacity=capacity)::ArenaCursor{T,D}
+end
 
 """Reset the cursor stack to start a preorder traversal at `root`."""
 @inline function reset!(c::ArenaCursor{T,D}, root::Int32) where {T,D}
@@ -691,11 +730,13 @@ end
 @inline reset!(c::ArenaCursor, root::ArenaNode) = reset!(c, root.idx)
 
 """Pop the next node in preorder (or return `nothing` when done)."""
-function next!(c::ArenaCursor{T,D}) where {T,D}
-    isempty(c.stack) && return nothing
+function next!(c::ArenaCursor{T,D})::Nullable{ArenaNode{T,D}} where {T,D}
+    if isempty(c.stack)
+        return Nullable(true, ArenaNode{T,D}(c.arena, Int32(0)))
+    end
 
     idx = pop!(c.stack)
-    n = ArenaNode{T,D}(c.arena, idx)
+    node = ArenaNode{T,D}(c.arena, idx)
 
     # Push children in reverse order so the leftmost child is visited next.
     d = @inbounds c.arena.degree[Int(idx)]
@@ -707,7 +748,7 @@ function next!(c::ArenaCursor{T,D}) where {T,D}
         end
     end
 
-    return n
+    return Nullable(false, node)
 end
 
 """Traverse a tree in preorder using a reusable cursor."""
@@ -717,9 +758,9 @@ function foreach_preorder!(f, root::ArenaNode{T,D}, cursor::ArenaCursor{T,D}) wh
 
     reset!(cursor, root)
     while true
-        n = next!(cursor)
-        n === nothing && break
-        f(n)
+        maybe_n = next!(cursor)
+        maybe_n.null && break
+        f(maybe_n[])
     end
     return nothing
 end

src/NodeUtils.jl

@@ -17,6 +17,8 @@ import ..NodeModule:
 import ..ValueInterfaceModule:
     pack_scalar_constants!, unpack_scalar_constants, count_scalar_constants, get_number_type
 
+import ..ArenaNodeModule: ArenaNode, ArenaCursor, reset!, next!
+
 """
     count_depth(tree::AbstractNode)::Int
 
@@ -70,7 +72,25 @@ has_operators(tree::AbstractExpressionNode) = tree.degree != 0
 Check if an expression is a constant numerical value, or
 whether it depends on input features.
 """
-is_constant(tree::AbstractExpressionNode) = all(t -> t.degree != 0 || t.constant, tree)
+is_constant(tree::AbstractExpressionNode) = !any(t -> t.degree == 0 && !t.constant, tree)
+
+# Specialized implementation for arena-backed nodes to keep DispatchDoctor inference concrete.
+function is_constant(tree::ArenaNode{T,D}) where {T,D}
+    cursor = ArenaCursor(tree; capacity=Int(getfield(tree, :idx)))
+    reset!(cursor, tree)
+    while true
+        maybe_n = next!(cursor)
+        maybe_n.null && break
+        n = maybe_n[]
+        arena = getfield(n, :arena)
+        i = Int(getfield(n, :idx))
+        d = @inbounds arena.degree[i]
+        if d == 0
+            @inbounds arena.constant[i] || return false
+        end
+    end
+    return true
+end
 
 """
     count_scalar_constants(tree::AbstractExpressionNode{T})::Int64 where {T}

test/test_arenanode.jl

@@ -117,7 +117,7 @@ const AN = DynamicExpressions.ArenaNodeModule
         @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
+        @test y_pf ≈ y_mut
 
         # Minimal rewrite prototype should preserve postfix validity:
         tree_constleft = 3.2 * x1
@@ -134,4 +134,82 @@ const AN = DynamicExpressions.ArenaNodeModule
         @test ok_after
         @test y_after ≈ y_before
     end
+
+    @testset "Arena allocations" begin
+        # DispatchDoctor checks in the test environment can dominate allocation counts.
+        # Measure these low-level allocation properties in a fresh process using the
+        # package project (dispatch doctor disabled by default there).
+        project_root = normpath(joinpath(@__DIR__, ".."))
+
+        alloc_script = raw"""
+            local_prefs = joinpath(dirname(Base.active_project()), "LocalPreferences.toml")
+            prefs_text = string(
+                "[DynamicExpressions]\n",
+                "dispatch_doctor_mode = ",
+                repr("disable"),
+                "\n",
+            )
+            write(local_prefs, prefs_text)
+            atexit(() -> rm(local_prefs; force=true))
+
+            using DynamicExpressions
+            const AN = DynamicExpressions.ArenaNodeModule
+
+            operators = OperatorEnum(1 => (sin, cos), 2 => (+, *))
+            x1 = DynamicExpressions.Node{Float64}(; feature=1)
+
+            function alloc_push_constant!(arena)
+                AN.push_constant!(arena, 1.0)
+                return nothing
+            end
+
+            function alloc_set_child!(parent, child)
+                set_child!(parent, child, 1)
+                return nothing
+            end
+
+            function alloc_copy_tree!(arena, tree)
+                AN._copy_to_arena!(arena, tree)
+                return nothing
+            end
+
+            arena_push = AN.Arena{Float64,2}(; capacity=16)
+
+            base_tree = sin(x1)
+            parent_arena = AN.Arena{Float64,2}(; capacity=16)
+            parent_idx = AN._copy_to_arena!(parent_arena, base_tree)
+            parent = AN.ArenaNode(parent_arena, parent_idx)
+
+            child_tree = x1 * 3.2
+            child_arena = AN.Arena{Float64,2}(; capacity=16)
+            child_idx = AN._copy_to_arena!(child_arena, child_tree)
+            child = AN.ArenaNode(child_arena, child_idx)
+
+            tree_large = sin(x1) + x1 * 3.2 + cos(x1)
+            arena_large = AN.Arena{Float64,2}(; capacity=64)
+
+            alloc_push_constant!(arena_push) # warmup
+            alloc_set_child!(parent, child) # warmup
+            alloc_copy_tree!(arena_large, tree_large) # warmup
+
+            println("push_constant=$(@allocated alloc_push_constant!(arena_push))")
+            println("set_child=$(@allocated alloc_set_child!(parent, child))")
+            println("copy_tree=$(@allocated alloc_copy_tree!(arena_large, tree_large))")
+        """
+
+        julia_bin = joinpath(Sys.BINDIR, Base.julia_exename())
+        cmd = `$(julia_bin) --startup-file=no --project=$(project_root) -e $(alloc_script)`
+        out = read(cmd, String)
+
+        allocs = Dict{String,Int}()
+        for m in eachmatch(r"(push_constant|set_child|copy_tree)=(\d+)", out)
+            allocs[m.captures[1]] = parse(Int, m.captures[2])
+        end
+
+        @test all(k -> haskey(allocs, k), ("push_constant", "set_child", "copy_tree"))
+        @test allocs["push_constant"] <= 1024
+        @test allocs["set_child"] <= 1024
+        @test allocs["copy_tree"] <= 1024
+    end
+
 end