ArenaNode for array-backed expressions - #167
Open
MilesCranmerBot wants to merge 79 commits into
Open
MilesCranmerBot wants to merge 79 commits into
MilesCranmerBot wants to merge 79 commits into
Conversation
Contributor
Benchmark Results (Julia v1)Time benchmarks
Memory benchmarks
|
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
February 21, 2026 12:56
e8acc61 to
847c2c6
Compare
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
June 9, 2026 09:23
847c2c6 to
04e763d
Compare
Member
|
@MilesCranmerBot comments from Claude... pls integrate. Maybe benchmark against master too, to see if it works... |
Contributor
Author
|
Integrated in What changed:
Validation: Benchmark ( |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #167 +/- ##
==========================================
+ Coverage 75.89% 84.70% +8.81%
==========================================
Files 30 31 +1
Lines 2912 3243 +331
==========================================
+ Hits 2210 2747 +537
+ Misses 702 496 -206 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please address above comments.
Also, shouldn't we have a tree_mapreduce? Then we wouldn't need to write a custom is_constant and so forth.
Btw, please expand your test suite. You should actually be using the full Node Interface tests. ArenaNode should pass ALL of them.
Contributor
Author
|
Addressed the review in the latest pushes ( Changes:
Local validation:
All addressed review threads are resolved. CI is pending on the latest push. |
MilesCranmer
requested changes
Jun 10, 2026
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
See new comments.
Also.
Please test with expressions types too. eg Expression with ArenaNode - do the full expression interface tests.
Also maybe try derivatives too?
MilesCranmer
requested changes
Jun 11, 2026
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please see comments above
Also please resolve any comments that have already been addressed.
Contributor
Author
|
Addressed in Changes:
Validation: Resolved the review threads addressed by this push. CI is pending on the latest commit. |
Member
|
@MilesCranmerBot there are still unresolved comments left. Also the CI is failing. Please address everything. |
Contributor
Author
|
Addressed in Changes:
Validation:
|
MilesCranmer
requested changes
Jun 11, 2026
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please address
Contributor
Author
|
Addressed in Changes:
Validation:
Resolved the review threads addressed by this push. CI is pending on the new commit. |
Member
|
@MilesCranmerBot downgrade compat is failing. Could you make a separate PR that fixes the lower bound? Or, if necessary, in this one (if there's a new dependency for example). Also why does it say only 0.54% of diff is actually tested... That is crazy. Should be 100%... |
Member
|
Also please resolve individual comments if already addressed. |
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
June 11, 2026 15:40
d401e65 to
92d4eef
Compare
Contributor
Author
|
Done.
Root cause: downgrade compat could select On the 0.54% Codecov number: that report is stale. The Codecov comment compares base |
Contributor
Author
|
Follow-up correction: the first pushed bound was too narrow for the Julia 1.10 downgrade path. I pushed LoweredCodeUtils = "2, 3.4 - 3.5.99"That keeps Julia 1.10 on the JET 0.9 / LoweredCodeUtils 2.x path, while excluding the failing LoweredCodeUtils 3.0.x path on Julia 1.12. Local fresh resolver/import checks:
|
Member
|
@MilesCranmerBot CI is still failing, pls fix |
Member
|
@MilesCranmerBot no excluding coverage files... You should try to actually improve the coverage as needed. |
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please address comments.
Contributor
Author
|
Fixed and pushed to the PR branch. Changes:
Verified locally with the full package test command plus targeted SymbolicUtils/Optim/LoopVectorization smoke tests. Remote checks are now green: CI matrix, downgrade compat, format, bench, and Codecov patch/project all pass. |
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot pls address comments
MilesCranmer
changed the title
MilesCranmer
changed the title
Member
|
@MilesCranmerBot can you please apply the following diff? Details 5 files changed, 367 insertions(+), 51 deletions(-)
diff --git a/src/ArenaNode.jl b/src/ArenaNode.jl
index 07e606e..faeafc4 100644
--- a/src/ArenaNode.jl
+++ b/src/ArenaNode.jl
@@ -9,32 +9,68 @@ import ..NodeModule:
unsafe_get_children,
get_child,
set_child!,
- set_children!
-
-"""Array-backed arena storing the fields of a tree node in a struct-of-arrays form.
+ set_children!,
+ count_nodes,
+ copy_node
+import ..NodeUtilsModule:
+ count_constant_nodes,
+ count_scalar_constants,
+ has_constants,
+ get_scalar_constants,
+ set_scalar_constants!
+import ..NodePreallocationModule: allocate_container, copy_into!
+import ..ValueInterfaceModule: get_number_type
+
+"""All per-node fields packed into a single isbits struct.
+
+Storing nodes as one `Vector{ArenaEntry}` (array-of-structs) makes whole-tree
+operations flat array operations: `copy` is a single `memcpy`, and traversals
+touch one contiguous stream of memory.
Indices are `Int32` and are 1-based. A child index of `0` indicates an empty slot.
+"""
+struct ArenaEntry{T,D}
+ val::T
+ children::NTuple{D,Int32}
+ feature::UInt16
+ degree::UInt8
+ op::UInt8
+ constant::Bool
+end
+
+@inline function _replace(
+ e::ArenaEntry{T,D};
+ val=e.val,
+ children=e.children,
+ feature=e.feature,
+ degree=e.degree,
+ op=e.op,
+ constant=e.constant,
+) where {T,D}
+ return ArenaEntry{T,D}(val, children, feature, degree, op, constant)
+end
+
+"""Array-backed arena storing the nodes of a tree contiguously.
This is an *experimental prototype* intended to provide an arena-backed representation
with a `Node`-like facade (`ArenaNode`) that supports existing tree algorithms that are
written against `AbstractExpressionNode`.
+
+The `compact` flag tracks whether `nodes` is exactly one postfix-ordered tree
+(children stored before parents, root last, no orphaned nodes and no shared
+subtrees). Trees built via `convert`/`copy` are compact; structural mutations
+through the facade may clear the flag, in which case whole-tree operations fall
+back to generic traversals. A structural `copy` re-compacts.
"""
struct Arena{T,D}
- degree::Vector{UInt8}
- constant::Vector{Bool}
- val::Vector{T}
- feature::Vector{UInt16}
- op::Vector{UInt8}
- children::Vector{NTuple{D,Int32}}
+ nodes::Vector{ArenaEntry{T,D}}
+ compact::Base.RefValue{Bool}
function Arena{T,D}(; capacity::Integer=0) where {T,D}
- degree = sizehint!(UInt8[], capacity)
- constant = sizehint!(Bool[], capacity)
- val = sizehint!(T[], capacity)
- feature = sizehint!(UInt16[], capacity)
- op = sizehint!(UInt8[], capacity)
- children = sizehint!(NTuple{D,Int32}[], capacity)
- return new{T,D}(degree, constant, val, feature, op, children)
+ return new{T,D}(sizehint!(ArenaEntry{T,D}[], capacity), Ref(true))
+ end
+ function Arena{T,D}(nodes::Vector{ArenaEntry{T,D}}, compact::Bool) where {T,D}
+ return new{T,D}(nodes, Ref(compact))
end
end
@@ -54,6 +90,13 @@ end
@inline ArenaNode(arena::Arena{T,D}, idx::Int32) where {T,D} = ArenaNode{T,D}(arena, idx)
+"""Whether `tree` is the root of a compact arena, so that the arena contents
+*are* the tree and whole-tree operations can act on the flat array directly."""
+@inline function is_compact_root(tree::ArenaNode)
+ a = getfield(tree, :arena)
+ return a.compact[] && Int(getfield(tree, :idx)) == length(a.nodes)
+end
+
@inline function _zero_children(::Val{D}) where {D}
return ntuple(_ -> Int32(0), Val(D))
end
@@ -67,13 +110,8 @@ end
op::UInt8,
children::NTuple{D,Int32},
) where {T,D}
- push!(arena.degree, degree)
- push!(arena.constant, constant)
- push!(arena.val, val)
- push!(arena.feature, feature)
- push!(arena.op, op)
- push!(arena.children, children)
- return Int32(length(arena.degree))
+ push!(arena.nodes, ArenaEntry{T,D}(val, children, feature, degree, op, constant))
+ return Int32(length(arena.nodes))
end
@inline function push_constant!(arena::Arena{T,D}, value) where {T,D}
@@ -117,15 +155,15 @@ Base.@constprop :aggressive @inline function Base.getproperty(
elseif k === :idx
return getfield(n, :idx)
elseif k === :degree
- return @inbounds getfield(n, :arena).degree[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].degree
elseif k === :constant
- return @inbounds getfield(n, :arena).constant[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].constant
elseif k === :val
- return @inbounds getfield(n, :arena).val[getfield(n, :idx)]::T
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].val::T
elseif k === :feature
- return @inbounds getfield(n, :arena).feature[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].feature
elseif k === :op
- return @inbounds getfield(n, :arena).op[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].op
elseif k === :children
return unsafe_get_children(n)
elseif k === :l
@@ -138,21 +176,26 @@ Base.@constprop :aggressive @inline function Base.getproperty(
end
@inline function Base.setproperty!(n::ArenaNode{T,D}, k::Symbol, v) where {T,D}
+ a = n.arena
i = n.idx
+ e = @inbounds a.nodes[i]
if k === :degree
- @inbounds n.arena.degree[i] = UInt8(v)
+ # Changing arity orphans or exposes child slots, so the flat layout can
+ # no longer be assumed to be exactly this tree.
+ UInt8(v) == e.degree || (a.compact[] = false)
+ @inbounds a.nodes[i] = _replace(e; degree=UInt8(v))
return v
elseif k === :constant
- @inbounds n.arena.constant[i] = Bool(v)
+ @inbounds a.nodes[i] = _replace(e; constant=Bool(v))
return v
elseif k === :val
- @inbounds n.arena.val[i] = convert(T, v)
+ @inbounds a.nodes[i] = _replace(e; val=convert(T, v))
return v
elseif k === :feature
- @inbounds n.arena.feature[i] = UInt16(v)
+ @inbounds a.nodes[i] = _replace(e; feature=UInt16(v))
return v
elseif k === :op
- @inbounds n.arena.op[i] = UInt8(v)
+ @inbounds a.nodes[i] = _replace(e; op=UInt8(v))
return v
elseif k === :l
set_child!(n, v, 1)
@@ -179,13 +222,13 @@ accessing them throws an `UndefRefError`.
"""
@generated function unsafe_get_children(n::ArenaNode{T,D}) where {T,D}
quote
- children = @inbounds getfield(n, :arena).children[getfield(n, :idx)]
+ children = @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].children
return Base.Cartesian.@ntuple($D, j -> _nullable_child(n, children[j]))
end
end
@inline function get_child(n::ArenaNode{T,D}, i::Int) where {T,D}
- c = @inbounds n.arena.children[n.idx][i]
+ c = @inbounds n.arena.nodes[n.idx].children[i]
c == 0 && throw(UndefRefError())
return ArenaNode(n.arena, c)
end
@@ -205,9 +248,14 @@ end
_copy_to_arena!(n.arena, child)
end
- old = @inbounds n.arena.children[n.idx]
- @inbounds n.arena.children[n.idx] = Base.setindex(old, idx, i)
- return ArenaNode(n.arena, idx)
+ a = n.arena
+ e = @inbounds a.nodes[n.idx]
+ if @inbounds(e.children[i]) != idx
+ # Relinking orphans the old child subtree (and may introduce sharing).
+ a.compact[] = false
+ @inbounds a.nodes[n.idx] = _replace(e; children=Base.setindex(e.children, idx, i))
+ end
+ return ArenaNode(a, idx)
end
@inline function set_children!(
@@ -236,17 +284,67 @@ end
idxs = Base.setindex(idxs, idx, i)
end
- @inbounds n.arena.children[n.idx] = idxs
+ a = n.arena
+ e = @inbounds a.nodes[n.idx]
+ if e.children != idxs
+ a.compact[] = false
+ @inbounds a.nodes[n.idx] = _replace(e; children=idxs)
+ end
return nothing
end
-"""Copy a tree into a new arena and return the new root node."""
-function Base.copy(tree::ArenaNode{T,D}; break_sharing::Val{BS}=Val(false)) where {T,D,BS}
+"""Copy a tree into a new arena and return the new root node.
+
+When `tree` is the root of a compact arena, this is a single flat copy of the
+node array (child indices are arena-relative, so they remain valid verbatim).
+Otherwise it falls back to a structural copy, which also re-compacts the
+resulting arena.
+
+This overloads `copy_node` (rather than `Base.copy`) since it is the generic
+entry point: `Base.copy(::AbstractExpressionNode)` forwards here, and the
+fallback `copy_node` would otherwise build a fresh arena per copied node via
+`constructorof`.
+"""
+function copy_node(tree::ArenaNode{T,D}; break_sharing::Val{BS}=Val(false)) where {T,D,BS}
+ if is_compact_root(tree)
+ return ArenaNode{T,D}(Arena{T,D}(copy(tree.arena.nodes), true), tree.idx)
+ end
arena = Arena{T,D}(; capacity=length(tree; break_sharing=Val(true)))
idx = _copy_to_arena!(arena, tree)
return ArenaNode{T,D}(arena, idx)
end
+"""Preallocate an arena for [`copy_into!`](@ref), enabling zero-allocation copies."""
+function allocate_container(
+ prototype::ArenaNode{T,D}, n::Union{Nothing,Integer}=nothing
+) where {T,D}
+ return Arena{T,D}(; capacity=@something(n, length(prototype)))
+end
+
+"""Copy `src` into the preallocated arena `dest`, reusing its storage.
+
+This is the steady-state copy path for population-based search: no allocations
+once `dest` has sufficient capacity.
+"""
+function copy_into!(
+ dest::Arena{T,D},
+ src::ArenaNode{T,D};
+ ref::Union{Nothing,Base.RefValue{<:Integer}}=nothing,
+) where {T,D}
+ @assert dest !== src.arena
+ if is_compact_root(src)
+ nodes = src.arena.nodes
+ resize!(dest.nodes, length(nodes))
+ copyto!(dest.nodes, nodes)
+ dest.compact[] = true
+ return ArenaNode{T,D}(dest, src.idx)
+ end
+ empty!(dest.nodes)
+ idx = _copy_to_arena!(dest, src)
+ dest.compact[] = true
+ return ArenaNode{T,D}(dest, idx)
+end
+
function _copy_to_arena!(arena::Arena{T,D}, tree::AbstractExpressionNode{T,D}) where {T,D}
d = tree.degree
if d == 0
@@ -266,7 +364,7 @@ end
"""Convert an existing tree into an arena-backed representation.
-This copies the entire tree into a fresh arena.
+This copies the entire tree into a fresh arena, in postfix (children-first) order.
"""
@inline function Base.convert(
::Type{ArenaNode{T,D}}, tree::AbstractExpressionNode{T,D}
@@ -281,6 +379,94 @@ end
return convert(ArenaNode{T,D}, tree)
end
+################################################################################
+# Flat whole-tree operations
+#
+# For a compact arena the node array *is* the tree, so tree-wide reductions
+# become linear array scans with no pointer chasing. Each of these falls back
+# to the generic traversal-based implementation when the invariant doesn't hold.
+################################################################################
+
+function count_nodes(tree::ArenaNode; break_sharing::Val{BS}=Val(false)) where {BS}
+ if is_compact_root(tree)
+ return length(tree.arena.nodes)
+ end
+ return invoke(count_nodes, Tuple{AbstractNode}, tree; break_sharing=Val(BS))::Int64
+end
+
+function count_constant_nodes(tree::ArenaNode)
+ if is_compact_root(tree)
+ return count(e -> e.degree == 0x00 && e.constant, tree.arena.nodes)
+ end
+ return invoke(count_constant_nodes, Tuple{AbstractExpressionNode}, tree)
+end
+
+function has_constants(tree::ArenaNode)
+ if is_compact_root(tree)
+ return any(e -> e.degree == 0x00 && e.constant, tree.arena.nodes)
+ end
+ return invoke(has_constants, Tuple{AbstractNode}, tree)
+end
+
+function count_scalar_constants(tree::ArenaNode{T}) where {T<:Number}
+ if is_compact_root(tree)
+ # For scalar `T`, each constant leaf stores exactly one scalar.
+ return count(e -> e.degree == 0x00 && e.constant, tree.arena.nodes)
+ end
+ return invoke(count_scalar_constants, Tuple{AbstractExpressionNode{T}}, tree)
+end
+
+"""Used by `NodeSampler` (random node selection in mutations), once per sample."""
+function Base.count(
+ f::F, tree::ArenaNode{T,D}; init=0, break_sharing::Val{BS}=Val(false)
+) where {F<:Function,T,D,BS}
+ if is_compact_root(tree)
+ a = tree.arena
+ c = init
+ @inbounds for i in 1:length(a.nodes)
+ c += f(ArenaNode{T,D}(a, Int32(i))) ? 1 : 0
+ end
+ return c
+ end
+ return invoke(Base.count, Tuple{F,AbstractNode}, f, tree; init, break_sharing=Val(BS))
+end
+
+"""For compact arenas, constants are gathered by a linear scan, and the
+returned `refs` are plain arena indices (which also remain valid in flat
+copies of the tree)."""
+function get_scalar_constants(
+ tree::ArenaNode{T}, ::Type{BT}=get_number_type(T)
+) where {T<:Number,BT}
+ if is_compact_root(tree)
+ nodes = tree.arena.nodes
+ n_constants = count(e -> e.degree == 0x00 && e.constant, nodes)
+ vals = Vector{T}(undef, n_constants)
+ refs = Vector{Int32}(undef, n_constants)
+ j = 0
+ @inbounds for i in eachindex(nodes)
+ e = nodes[i]
+ if e.degree == 0x00 && e.constant
+ j += 1
+ vals[j] = e.val
+ refs[j] = Int32(i)
+ end
+ end
+ return vals, refs
+ end
+ return invoke(get_scalar_constants, Tuple{AbstractExpressionNode{T},Type{BT}}, tree, BT)
+end
+
+function set_scalar_constants!(
+ tree::ArenaNode{T}, constants, refs::AbstractVector{Int32}
+) where {T<:Number}
+ nodes = tree.arena.nodes
+ @inbounds for j in eachindex(refs, constants)
+ i = refs[j]
+ nodes[i] = _replace(nodes[i]; val=constants[j]::T)
+ end
+ return nothing
+end
+
################################################################################
# Cursor + reusable stack (prototype)
################################################################################
@@ -325,11 +511,10 @@ function next!(c::ArenaCursor{T,D})::Nullable{ArenaNode{T,D}} where {T,D}
node = ArenaNode{T,D}(c.arena, idx)
# Push children in reverse order so the leftmost child is visited next.
- d = @inbounds c.arena.degree[idx]
- if d != 0
- child_idxs = @inbounds c.arena.children[idx]
- @inbounds for i in d:-1:1
- child = child_idxs[i]
+ e = @inbounds c.arena.nodes[idx]
+ if e.degree != 0
+ @inbounds for i in e.degree:-1:1
+ child = e.children[i]
child != 0 && push!(c.stack, child)
end
end
diff --git a/src/DynamicExpressions.jl b/src/DynamicExpressions.jl
index b1fa6c6..5e2688f 100644
--- a/src/DynamicExpressions.jl
+++ b/src/DynamicExpressions.jl
@@ -8,9 +8,9 @@ using DispatchDoctor: @stable, @unstable
include("ExtensionInterface.jl")
include("OperatorEnum.jl")
include("Node.jl")
- include("ArenaNode.jl")
include("NodeUtils.jl")
include("NodePreallocation.jl")
+ include("ArenaNode.jl")
include("Strings.jl")
include("Evaluate.jl")
include("EvaluateDerivative.jl")
diff --git a/test/runtests.jl b/test/runtests.jl
index fa16d57..7bad778 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -56,7 +56,7 @@ if "jet" in test_names
)
s_mod = string(mod.mod)
any(report.vst) do vst
- occursin(s_mod, string(JET.linfomod(vst.linfo)))
+ return occursin(s_mod, string(JET.linfomod(vst.linfo)))
end
end
# On JET 0.10, `target_defined_modules` is not available and also
@@ -82,6 +82,10 @@ testitem_suffixes = String[]
if "main" in test_names
push!(testitem_suffixes, joinpath("test", "unittest.jl"))
push!(testitem_suffixes, joinpath("test", "test_optim.jl"))
+ # NOTE: `@testitem`s defined in a file that `unittest.jl` merely `include`s are
+ # attributed to *their own* filename by TestItemRunner, so files with their own
+ # testitems must be listed here explicitly or they will be silently skipped.
+ push!(testitem_suffixes, joinpath("test", "test_arenanode.jl"))
end
if "optim" in test_names
push!(testitem_suffixes, joinpath("test", "test_optim.jl"))
diff --git a/test/test_arenanode.jl b/test/test_arenanode.jl
index 98f0fc7..342e436 100644
--- a/test/test_arenanode.jl
+++ b/test/test_arenanode.jl
@@ -118,7 +118,8 @@ end
@test string_tree(rewritten, operators) == "sin(x1 * 2.0)"
bad_children = (
- DynamicExpressions.Nullable{Node{Float64,2}}(true), Node{Float32}(; val=1.0f0)
+ DynamicExpressions.Nullable(true, Node{Float64}(; val=0.0)),
+ Node{Float32}(; val=1.0f0),
)
@test_throws ArgumentError set_children!(rewritten, bad_children)
@@ -203,7 +204,7 @@ end
@test grad3[1, :] ≈ fill(1.0, 5)
d_ex = gradient(AutoZygote(), expr_const) do ex
- sum(ex(ones(1, 5)))
+ return sum(ex(ones(1, 5)))
end
@test extract_gradient(d_ex, expr_const) ≈ [5.0]
end
@@ -232,3 +233,119 @@ end
end
end
end
+
+@testitem "ArenaNode flat copy and whole-tree fast paths" begin
+ using DynamicExpressions
+ using DynamicExpressions: Node, copy_node
+ using DynamicExpressions.NodePreallocationModule: allocate_container, copy_into!
+
+ const AN = DynamicExpressions.ArenaNodeModule
+
+ operators = OperatorEnum(; binary_operators=[+, -, *, /], unary_operators=[sin, cos])
+ x1 = Node{Float64}(; feature=1)
+ x2 = Node{Float64}(; feature=2)
+ tree = sin(x1 * 3.2 - 0.9) + x2 * (x1 - 0.5)
+ atree = convert(AN.ArenaNode{Float64}, tree)
+
+ @testset "compact flat copy" begin
+ @test AN.is_compact_root(atree)
+ c = copy(atree)
+ @test c.arena !== atree.arena
+ @test convert(Node, c) == tree
+ c.l.l.r.val = 99.0
+ @test convert(Node, atree) == tree
+ @test convert(Node, c) != tree
+ end
+
+ @testset "subtree copy falls back and re-compacts" begin
+ sub = atree.l
+ @test !AN.is_compact_root(sub)
+ csub = copy(sub)
+ @test AN.is_compact_root(csub)
+ @test convert(Node, csub) == tree.l
+ end
+
+ @testset "structural mutation invalidates fast paths" begin
+ mutated = convert(AN.ArenaNode{Float64}, tree)
+ set_child!(mutated, convert(AN.ArenaNode{Float64}, cos(x2)), 2)
+ @test !mutated.arena.compact[]
+ expected = copy(tree)
+ set_child!(expected, cos(x2), 2)
+ @test convert(Node, mutated) == expected
+ @test count_nodes(mutated) == count_nodes(expected)
+ recompacted = copy(mutated)
+ @test AN.is_compact_root(recompacted)
+ @test count_nodes(recompacted) == count_nodes(expected)
+
+ leafed = convert(AN.ArenaNode{Float64}, tree)
+ node = leafed.r
+ node.degree = 0
+ node.constant = true
+ node.val = 1.0
+ @test !leafed.arena.compact[]
+ expected2 = copy(tree)
+ expected2.r = Node{Float64}(; val=1.0)
+ @test convert(Node, leafed) == expected2
+ @test count_nodes(leafed) == count_nodes(expected2)
+ end
+
+ @testset "preallocated copy_into!" begin
+ dest = allocate_container(atree)
+ out = copy_into!(dest, atree)
+ @test out.arena === dest
+ @test convert(Node, out) == tree
+ out2 = copy_into!(dest, atree)
+ @test convert(Node, out2) == tree
+ end
+
+ @testset "copy_node entry point" begin
+ c = copy_node(atree)
+ @test c.arena !== atree.arena
+ @test convert(Node, c) == tree
+ end
+
+ @testset "Expression-level preallocated copy (SR mutation path)" begin
+ ex = Expression(
+ convert(AN.ArenaNode{Float64}, tree);
+ operators=operators,
+ variable_names=["x1", "x2"],
+ )
+ container = allocate_container(ex)
+ ex2 = copy_into!(container, ex)
+ @test convert(Node, DynamicExpressions.get_tree(ex2)) == tree
+ end
+
+ @testset "whole-tree scans match Node" begin
+ @test count_nodes(atree) == count_nodes(tree)
+ @test count(t -> t.degree == 2, atree) == count(t -> t.degree == 2, tree)
+ @test count(t -> t.degree == 0, atree.l) == count(t -> t.degree == 0, tree.l)
+ @test length(atree) == length(tree)
+ @test count_constant_nodes(atree) == count_constant_nodes(tree)
+ @test has_constants(atree) == has_constants(tree)
+ leaf = convert(AN.ArenaNode{Float64}, Node{Float64}(; feature=1))
+ @test !has_constants(leaf)
+ @test count_constant_nodes(leaf) == 0
+ end
+
+ @testset "scalar constants via arena indices" begin
+ fresh = convert(AN.ArenaNode{Float64}, tree)
+ vals, refs = get_scalar_constants(fresh)
+ @test refs isa Vector{Int32}
+ @test DynamicExpressions.count_scalar_constants(fresh) == length(vals)
+ @test vals == first(get_scalar_constants(tree))
+ set_scalar_constants!(fresh, vals .* 2, refs)
+ @test first(get_scalar_constants(fresh)) == vals .* 2
+
+ # Indices remain valid in flat copies of the tree:
+ c = copy(fresh)
+ set_scalar_constants!(c, vals, refs)
+ @test first(get_scalar_constants(c)) == vals
+
+ # Non-compact trees fall back to the generic Ref-based path:
+ sub = fresh.l
+ vsub, rsub = get_scalar_constants(sub)
+ @test vsub == first(get_scalar_constants(convert(Node, sub)))
+ set_scalar_constants!(sub, vsub .+ 1, rsub)
+ @test first(get_scalar_constants(sub)) == vsub .+ 1
+ end
+end
diff --git a/test/test_arenanode_allocations.jl b/test/test_arenanode_allocations.jl
index fc12b93..429de7b 100644
--- a/test/test_arenanode_allocations.jl
+++ b/test/test_arenanode_allocations.jl
@@ -1,5 +1,6 @@
using Test
using DynamicExpressions
+using DynamicExpressions.NodePreallocationModule: allocate_container, copy_into!
const AN = DynamicExpressions.ArenaNodeModule
@@ -26,6 +27,11 @@ function alloc_eval_tree(tree, X, operators)
return nothing
end
+function alloc_copy_into!(dest, tree)
+ copy_into!(dest, tree)
+ return nothing
+end
+
arena_push = AN.Arena{Float64,2}(; capacity=128)
base_tree = sin(x1)
@@ -40,6 +46,7 @@ child = AN.ArenaNode(child_arena, child_idx)
tree_large = sin(x1) + x1 * 3.2 + cos(x1)
atree_large = convert(AN.ArenaNode{Float64}, tree_large)
+copy_dest = allocate_container(atree_large)
arena_large = AN.Arena{Float64,2}(; capacity=128)
X = randn(Float64, 1, 1_000)
@@ -49,12 +56,14 @@ for _ in 1:5
alloc_copy_tree!(arena_large, tree_large)
alloc_eval_tree(tree_large, X, operators)
alloc_eval_tree(atree_large, X, operators)
+ alloc_copy_into!(copy_dest, atree_large)
end
alloc_counts = Dict(
"push_constant" => @allocations(alloc_push_constant!(arena_push)),
"set_child" => @allocations(alloc_set_child!(parent, child)),
"copy_tree" => @allocations(alloc_copy_tree!(arena_large, tree_large)),
+ "copy_into" => @allocations(alloc_copy_into!(copy_dest, atree_large)),
)
alloc_bytes = Dict(
"eval_node" => @allocated(alloc_eval_tree(tree_large, X, operators)),
@@ -64,4 +73,5 @@ alloc_bytes = Dict(
@test alloc_counts["push_constant"] == 0
@test alloc_counts["set_child"] == 0
@test alloc_counts["copy_tree"] == 0
+@test alloc_counts["copy_into"] == 0
@test alloc_bytes["eval_arena"] <= max(1024, ceil(Int, 1.10 * alloc_bytes["eval_node"]))
-- |
Documents the self-buffering experiment result: the generic fused unbuffered evaluator benches even with a plan running on a self-allocated pool, while allocating half the bytes, so the plan stays buffer-gated. Also shows buffered Node is slower than unbuffered Node (strided buffer rows), which is the headroom the plan's contiguous slots reclaim.
ArenaNode.jl keeps the types and the AbstractExpressionNode interface (~510 lines); ArenaNodeEval.jl holds the buffered plan evaluator (~530 lines), included from within the module. Also drops a tutorial-grade comment and records why the plan does not self-allocate a pool.
…t compacting copy
CI:
- _scalar_degn binds T via Tuple{T,Vararg{T}} and derives the arity from
the tuple type, fixing Aqua's unbound-args failure
- formatted with JuliaFormatter v1 (the CI pin; local runs had used v2)
Investigation of the slow any(f, ::ArenaNode) found the generic
AbstractNode traversal machinery is already fast on ArenaNode (within 15%
of Node, zero allocations) -- the custom _arena_any/_arena_mapreduce/
is_constant overrides were unnecessary, and the closure-based recursion in
the any override was nondeterministically up to 8x slower (inlining of the
Base.any layers is compilation-order dependent). So:
- the traversal overrides are deleted; ArenaNode rides the generic
machinery, except a 3-line compact-arena flat scan for (1.4-1.7x
faster than Node, and has_constants/is_constant inherit it)
- the unset-child UndefRefError guard moves from the deleted override into
_load_entry at the facade layer, so poison facades throw like Node's
undefined fields for every generic consumer
Copying out of a non-compact arena is now an entry-level compacting write
(_write_subtree! into a pre-sized vector; no facade traversal, no per-push
growth checks), also used by copy_into! and cross-arena attachment.
copy(::ArenaNode): compact 5.7-11x faster than Node at n=15-255;
non-compact crosses over near n=20 (1.1-1.2x above).
ArenaNodeEval.jl is now its own module included from DynamicExpressions.jl
rather than a nested include. The eval benchmark also reports unbuffered
paths.
Property tests (using the existing supposition_utils generators): Node -> ArenaNode -> Node round-trips with all read-only interface results equal; evaluation matches Node (unbuffered exactly; the buffered plan's ok flag is best-effort, values must agree whenever both report ok); random valid mutations applied to both representations keep them equivalent; copy re-compacts and preserves evaluation.
…_nodes hash went through the generic tree_mapreduce machinery (0.64x vs Node); direct recursion over entries through the shared leaf_hash/branch_hash combinators recovers it to 0.92x -- the rest is the hash arithmetic itself, which must be identical across representations since cross-representation == requires equal hashes (now pinned by the supposition round-trip property). count_constant_nodes on a compact arena is a flat count over the entry vector (1.6x vs Node).
The generic traversal skeleton re-reads the entry through the facade for every field access and shuffles Nullable-wrapped children, which taxed all tree_mapreduce clients at once (hash 0.64x, collect 0.70x, count_depth 0.76x, count_constant_nodes 0.73x vs Node). One override of the primitive -- read each entry exactly once, dispatch on the local degree, hand f the facade -- normalizes the whole family (1.0x for collect/count_depth/ count_constant_nodes, 0.94x index_constant_nodes, 0.85x hash) in both compact and non-compact arenas. This replaces the per-function fixes: the previous hash and count_constant_nodes specializations are deleted, and since the generic hash combinators now just run on the fast skeleton, hash equality across representations holds again (pinned in the supposition round-trip property).
At 31-node trees the generic any falls to 0.52x vs Node on non-compact arenas (facade re-reads per field, like the tree_mapreduce skeleton); the entry-level recursion recovers to 0.67x. The remainder is f reading fields through the facade (arena -> nodes -> entry per access vs Node's single deref), inherent to the facade contract; compact arenas bypass it entirely via the flat scan (1.8x ahead of Node at n=31).
A linear postfix scan over compact arenas benches ~20% faster for pure reductions (hash 1.03x vs Node), but f-application order (parent first, siblings left to right) is observable API through collect/filter/foreach: the postfix scan reorders them, applying positional mutations to different nodes than Node would. Caught by the supposition properties with shrunk counterexamples (abs(5.0e-324), one set_val). Recursion stays; the comment on _entry_mapreduce documents the constraint.
The flat scan was the recursive traversal with recursion disabled, so the
two share a single _entry_any whose children loop is gated on a static
Val{recurse} flag: the compact path drives it from a flat loop with the
recursion compiled out, the non-compact path recurses. Performance
unchanged (compact 1.6-2.0x vs Node at n=31, non-compact 0.63-0.68x).
…tant ops One driver for order-free whole-tree visits, with both walks side by side and the unused one compiled out via a static flag: compact arenas walk the entry array (indices valid by construction, so the load is uncheckable and dead when unused); non-compact arenas walk the children pointers with the per-degree dispatch unrolled like the generic any (a children loop benches ~40% slower; a worklist allocates and loses at SR sizes; closure-driven recursion through Base.any inlines nondeterministically -- all measured). any/count_nodes/get_scalar_constants are one-line clients, replacing the recursive _entry_any, the count_nodes invoke fallback, and the filter_map path. Also, on the canonical-form principle (pure or order-free ops on a compact arena may use the entries directly): - == between compact trees is a flat field-compare (equal trees have identical entry vectors); 2.7x vs Node, from 1.1x - count_constant_nodes: flat count when compact (2.0x), entry-level mapreduce otherwise - set_scalar_constants! writes entries directly: val-only writes cannot change structure, so the setindex! degree/children diff is skipped - @inline restored on getproperty (the removal had been validated against an eval-only benchmark that never exercises property reads) n=31 vs Node: compact 1.4-3.4x on ten ops, floor 0.81x (hash); non-compact 0.85-1.75x, floor 0.65x (copy_into!, the re-compaction price that wins back above n~20).
Older Supposition versions resolved by the downgrade-compat CI job lack length(::Data.OneOf); a single multi-arg map selecting the mutation kind by integer generates the same tuples on any version.
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
August 1, 2026 14:28
60d0023 to
9850b0b
Compare
Member
|
@MilesCranmerBot please merge or rebase on latest master |
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
September 18, 2026 23:36
2846c3b to
63939c2
Compare
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
JET 0.10 mis-analyzes package loading on Julia 1.13; allow JET 0.12, which added 1.13 support, and use the target_modules configuration that replaced the removed target_defined_modules. PrecompileTools 1.3.0-1.3.2 reference the renamed Base.StaticData, and TestItemRunner below 1.3.1 calls the removed Test.push_testset, so the downgrade job broke on Julia 1.13. The PrecompileTools range keeps 1.2.x for Julia 1.10 and 1.11. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
The isa-chain guard left buffer and cX as abstract unions under whole-package analysis, so JET 0.12 reported next_index!(::Nothing) and an unmatched _arena_eval call. Selecting the fast path by dispatch on the matrix buffer and matrix input makes both concrete. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Zygote 0.7.0 through 0.7.6 import typesof from InteractiveUtils, where Julia 1.13 no longer declares it, so precompilation fails. 0.7.7 switched to Base.typesof and still supports Julia 1.10. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Under the caller-owned ArrayBuffer contract the plan evaluator had to address the buffer row by row, which is 15-22% slower than running the generic evaluator on the arena facade (parity with Node within 2%). ArenaNode keeps its storage and copy advantages and evaluates through the generic path only. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
September 19, 2026 16:09
16b4904 to
6d06b9d
Compare
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
The generic constructorof builds `ArenaNode{T,D} where T` with an unbounded TypeVar. With the struct declared as `ArenaNode{T<:Number,D}` that UnionAll violates the bound, so Julia cannot cache it and every leaf construction allocated a fresh DataType and SimpleVector (256 B, ~2 us). Node has no bound and folds to a constant. SR append/insert mutations drop from 1.3 us to 0.6-0.7 us, matching Node.
Also count nodes of a non-compact arena with _entry_mapreduce instead of a boxed counter.
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Base grows small vectors by 3-4x on every reallocation, so splicing a subtree into a 20-entry arena allocated 70-86 entries. Arena now tracks the exact buffer length and doubles it itself, subtree appends are sized by the exact subtree count instead of the whole source length, and allocate_container reserves doubling headroom so the mutations applied after copy_into! do not reallocate. The compact flag becomes a plain field of the now-mutable struct. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Single-entry pushes keep geometric growth for amortized O(1) appends. Whole-subtree appends (crossover splicing into a fresh exact-size copy) grow to the exact size instead: the splice is a one-shot event and the headroom was never used. allocate_container leaves room for one operator node with its leaves rather than doubling. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
A copied tree is usually edited once before it is copied again (a mutation on a container copy, or the subtree a crossover splices in). Leaving room for one operator node with its leaves on every copy, as allocate_container already did, lets those edits append without reallocating the arena. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
…rena ArenaNode(; op, children) allocated a fresh one-entry arena and copied every child into it, growing the buffer once per child. Mutations that wrap an existing subtree (insert, prepend, append) therefore copied that subtree once more than Node does. The new node now lives in the arena of its first ArenaNode child, which attaches by reference as with Node; only the other children are copied, and they fit the edit headroom a copy already carries. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Prototype
ArenaNode: an array-backed arena representation for expressions, implementing the fullNodeInterface. Nodes live in oneVector{ArenaEntry}(array of structs), so a tree is a single contiguous block of memory:copyis one vector copy and traversals stream through one allocation.Measured on a 63-node
Float64tree (Julia 1.12.7, M1 Pro):NodeArenaNodeBase.summarysizecopyEvaluation goes through the generic evaluator via the
NodeInterfaceand runs at parity withNode(within 2%, buffered and unbuffered, trees of 7 to 31 nodes). An earlier revision carried a dedicated buffered plan evaluator; under the caller-ownedArrayBuffercontract it was 15-22% slower than the generic path, so it was dropped.