feat: better interface for children

Author: MilesCranmer

src/Evaluate.jl

@@ -3,7 +3,7 @@ module EvaluateModule
 using DispatchDoctor: @stable, @unstable
 
 import ..NodeModule:
-    AbstractExpressionNode, constructorof, max_degree, children, with_type_parameters
+    AbstractExpressionNode, constructorof, max_degree, get_children, with_type_parameters
 import ..StringsModule: string_tree
 import ..OperatorEnumModule: OperatorEnum, GenericOperatorEnum
 import ..UtilsModule: fill_similar, counttuple, ResultOk
@@ -343,7 +343,7 @@ end
 ) where {T,degree,OPS}
     nops = length(OPS.types[degree].types)
     return quote
-        cs = children(tree, Val($degree))
+        cs = get_children(tree, Val($degree))
         Base.Cartesian.@nexprs(
             $degree,
             i -> begin
@@ -727,7 +727,7 @@ end
 ) where {T,degree,OPS}
     nops = length(OPS.types[degree].types)
     get_inputs = quote
-        cs = children(tree, Val($degree))
+        cs = get_children(tree, Val($degree))
         Base.Cartesian.@nexprs(
             $degree,
             i -> begin

src/Interfaces.jl

@@ -12,7 +12,7 @@ using ..NodeModule:
     constructorof,
     default_allocator,
     with_type_parameters,
-    children,
+    get_children,
     leaf_copy,
     leaf_convert,
     leaf_hash,
@@ -226,12 +226,12 @@ function _check_create_node(tree::AbstractExpressionNode)
     NT = with_type_parameters(N, Float16)
     return NT() isa NT
 end
-function _check_children(tree::AbstractExpressionNode{T,D}) where {T,D}
+function _check_get_children(tree::AbstractExpressionNode{T,D}) where {T,D}
     tree.degree == 0 && return true
-    return children(tree) isa Tuple{typeof(tree),Vararg{typeof(tree)}} &&
-               children(tree, Val(D)) isa Tuple &&
-               length(children(tree, Val(D))) == D &&
-               length(children(tree, Val(1))) == 1
+    return get_children(tree) isa Tuple{typeof(tree),Vararg{typeof(tree)}} &&
+           get_children(tree, Val(D)) isa Tuple &&
+           length(get_children(tree, Val(D))) == D &&
+           length(get_children(tree, Val(1))) == 1
 end
 function _check_copy(tree::AbstractExpressionNode)
     return copy(tree) isa typeof(tree)
@@ -308,19 +308,19 @@ function _check_leaf_equal(tree::AbstractExpressionNode)
 end
 function _check_branch_copy(tree::AbstractExpressionNode)
     tree.degree == 0 && return true
-    return branch_copy(tree, children(tree, Val(tree.degree))...) isa typeof(tree)
+    return branch_copy(tree, get_children(tree, Val(tree.degree))...) isa typeof(tree)
 end
 function _check_branch_copy_into!(tree::AbstractExpressionNode{T}) where {T}
     tree.degree == 0 && return true
     new_branch = constructorof(typeof(tree))(; val=zero(T))
     ret = branch_copy_into!(
-        new_branch, tree, map(copy, children(tree, Val(tree.degree)))...
+        new_branch, tree, map(copy, get_children(tree, Val(tree.degree)))...
     )
     return new_branch == tree && ret === new_branch
 end
 function _check_branch_convert(tree::AbstractExpressionNode)
     tree.degree == 0 && return true
-    return branch_convert(typeof(tree), tree, children(tree, Val(tree.degree))...) isa
+    return branch_convert(typeof(tree), tree, get_children(tree, Val(tree.degree))...) isa
            typeof(tree)
 end
 function _check_branch_hash(tree::AbstractExpressionNode)
@@ -367,7 +367,7 @@ end
 ni_components = (
     mandatory = (
         create_node = "creates a new instance of the node type" => _check_create_node,
-        children = "returns the children of the node" => _check_children,
+        get_children = "returns the children of the node" => _check_get_children,
         copy = "returns a copy of the tree" => _check_copy,
         hash = "returns the hash of the tree" => _check_hash,
         any = "checks if any element of the tree satisfies a condition" => _check_any,

src/Node.jl

@@ -176,34 +176,56 @@ function get_poison(n::AbstractNode)
     return n
 end
 
+@inline function get_children(node::AbstractNode)
+    return getfield(node, :children)
+end
+@inline function get_children(node::AbstractNode, ::Val{n}) where {n}
+    cs = get_children(node)
+    return ntuple(i -> cs[i], Val(Int(n)))
+end
+@inline function get_child(n::AbstractNode{D}, i::Int) where {D}
+    return get_children(n)[i]
+end
+@inline function set_child!(n::AbstractNode{D}, child::AbstractNode{D}, i::Int) where {D}
+    set_children!(n, Base.setindex(get_children(n), child, i))
+    return child
+end
+@inline function set_children!(n::AbstractNode{D}, children::NTuple{D2,AbstractNode{D}}) where {D,D2}
+    if D === D2
+        n.children = children
+    else
+        poison = get_poison(n)
+        # We insert poison at the end of the tuple so that
+        # errors will appear loudly if accessed.
+        # This poison should be efficient to insert. So
+        # for simplicity, we can just use poison == n, which
+        # will trigger infinite recursion errors if accessed.
+        n.children = ntuple(i -> i <= D2 ? children[i] : poison, Val(D))
+    end
+end
+
 macro make_accessors(node_type)
     esc(quote
         @inline function Base.getproperty(n::$node_type, k::Symbol)
             if k == :l
                 # TODO: Should a depwarn be raised here? Or too slow?
-                return getfield(n, :children)[1]
+                return $(get_child)(n, 1)
             elseif k == :r
-                return getfield(n, :children)[2]
+                return $(get_child)(n, 2)
             else
                 return getfield(n, k)
             end
         end
         @inline function Base.setproperty!(n::$node_type, k::Symbol, v)
             if k == :l
                 if isdefined(n, :children)
-                    old = getfield(n, :children)
-                    setfield!(n, :children, (v, old[2]))
-                    v
+                    $(set_child!)(n, v, 1)
                 else
-                    poison = $(get_poison)(n)
-                    setfield!(n, :children, (v, poison))
+                    $(set_children!)(n, (v,))
                     v
                 end
             elseif k == :r
-                # TODO: Remove this assert once we know that this is safe
-                old = getfield(n, :children)
-                setfield!(n, :children, (old[1], v))
-                v
+                $(set_child!)(n, v, 2)
             else
                 T = fieldtype(typeof(n), k)
                 if v isa T
@@ -222,12 +244,6 @@ end
 @make_accessors GraphNode
 # TODO: Disable the `.l` accessors eventually, once the codebase is fully generic
 
-@inline children(node::AbstractNode) = node.children
-@inline function children(node::AbstractNode, ::Val{n}) where {n}
-    cs = children(node)
-    return ntuple(i -> cs[i], Val(Int(n)))
-end
-
 ################################################################################
 #! format: on
 
@@ -273,11 +289,11 @@ include("base.jl")
 @inline function (::Type{N})(
     ::Type{T1}=Undefined; val=nothing, feature=nothing, op=nothing, l=nothing, r=nothing, children=nothing, allocator::F=default_allocator,
 ) where {T1,N<:AbstractExpressionNode{T} where T,F}
-    _children = if l !== nothing && r === nothing
-        @assert children === nothing
+    _children = if !isnothing(l) && isnothing(r)
+        @assert isnothing(children)
         (l,)
-    elseif l !== nothing && r !== nothing
-        @assert children === nothing
+    elseif !isnothing(l) && !isnothing(r)
+        @assert isnothing(children)
         (l, r)
     else
         children
@@ -328,8 +344,7 @@ end
     n = allocator(N, T)
     n.degree = D2
     n.op = op
-    poison = get_poison(n)
-    n.children = ntuple(i -> i <= D2 ? convert(NT, children[i]) : poison, Val(max_degree(N)))
+    set_children!(n, children)
     return n
 end
 
@@ -398,7 +413,7 @@ function set_node!(tree::AbstractExpressionNode, new_tree::AbstractExpressionNod
         end
     else
         tree.op = new_tree.op
-        tree.children = new_tree.children
+        set_children!(tree, get_children(new_tree))
     end
     return nothing
 end

src/NodePreallocation.jl

@@ -7,7 +7,7 @@ using ..NodeModule:
     leaf_copy,
     branch_copy,
     set_node!,
-    get_poison
+    set_children!
 
 """
     allocate_container(prototype::AbstractExpressionNode, n=nothing)
@@ -60,8 +60,7 @@ function branch_copy_into!(
 ) where {T,D,N<:AbstractExpressionNode{T,D},M}
     dest.degree = M
     dest.op = src.op
-    poison = get_poison(dest)
-    dest.children = ntuple(i -> i <= M ? children[i] : poison, D)
+    set_children!(dest, children)
     return dest
 end
 

src/NodeUtils.jl

@@ -6,7 +6,7 @@ import ..NodeModule:
     Node,
     preserve_sharing,
     constructorof,
-    get_poison,
+    set_children!,
     copy_node,
     count_nodes,
     tree_mapreduce,
@@ -156,10 +156,9 @@ mutable struct NodeIndex{T,D} <: AbstractNode{D}
         ::Type{_T}, ::Val{_D}, child::NodeIndex{_T,_D}, childs::Vararg{NodeIndex{_T,_D},_D2}
     ) where {_T,_D,_D2}
         node = NodeIndex(_T, Val(_D))
-        poison = get_poison(node)
         children = (child, childs...)
         node.degree = _D2 + 1
-        node.children = ntuple(i -> i <= _D2 + 1 ? children[i] : poison, Val(_D))
+        set_children!(node, children)
         return node
     end
 end

src/ReadOnlyNode.jl

@@ -3,7 +3,7 @@ module ReadOnlyNodeModule
 using DispatchDoctor: @unstable
 
 using ..NodeModule: AbstractExpressionNode, Node
-import ..NodeModule: default_allocator, with_type_parameters, constructorof, children
+import ..NodeModule: default_allocator, with_type_parameters, constructorof, get_children
 
 abstract type AbstractReadOnlyNode{T,D,N<:AbstractExpressionNode{T,D},IS_REF} <:
               AbstractExpressionNode{T,D} end
@@ -38,8 +38,8 @@ Base.getindex(n::AbstractReadOnlyNode{T,D,N,true} where {T,D,N}) = n
         return out
     end
 end
-@inline function children(node::AbstractReadOnlyNode, ::Val{n}) where {n}
-    return map(ReadOnlyNode, children(inner(node), Val(n)))
+@inline function get_children(node::AbstractReadOnlyNode)
+    return map(ReadOnlyNode, get_children(inner(node)))
 end
 function Base.setproperty!(::AbstractReadOnlyNode, ::Symbol, v)
     return error("Cannot set properties on a ReadOnlyNode")

src/base.jl

@@ -137,7 +137,7 @@ end
                 Base.Cartesian.@nif(
                     $D,
                     i -> i == d,
-                    i -> let cs = children(tree, Val(i))
+                    i -> let cs = get_children(tree, Val(i))
                         Base.Cartesian.@ncall(
                             i,
                             mapreducer.op,
@@ -182,7 +182,7 @@ By using this instead of tree_mapreduce, we can take advantage of early exits.
 
         return (
             @inline(f(tree)) || Base.Cartesian.@nif(
-                $D, i -> deg == i, i -> let cs = children(tree, Val(i))
+                $D, i -> deg == i, i -> let cs = get_children(tree, Val(i))
                     Base.Cartesian.@nany(i, j -> any(f, cs[j]))
                 end
             )
@@ -226,9 +226,12 @@ end
                 branch_equal(a, b) && Base.Cartesian.@nif(
                     $D,
                     i -> deg == i,
-                    i -> let cs_a = children(a, Val(i)), cs_b = children(b, Val(i))
-                        Base.Cartesian.@nall(i, j -> inner_is_equal(cs_a[j], cs_b[j], id_maps))
-                    end
+                    i ->
+                        let cs_a = get_children(a, Val(i)), cs_b = get_children(b, Val(i))
+                            Base.Cartesian.@nall(
+                                i, j -> inner_is_equal(cs_a[j], cs_b[j], id_maps)
+                            )
+                        end
                 )
             )
         end

test/test_n_arity_nodes.jl

@@ -32,7 +32,7 @@
     @test n_bin.children[1] === n_bin_leaf1
     @test n_bin.children[2] === n_bin_leaf2
     @test n_bin.children[3] === n_bin # Poison
-    @test DynamicExpressions.NodeModule.children(n_bin, Val(2)) ==
+    @test DynamicExpressions.NodeModule.get_children(n_bin, Val(2)) ==
         (n_bin_leaf1, n_bin_leaf2)
     # .l and .r should work for Node{T,3} due to general @make_accessors Node
     @test n_bin.l === n_bin_leaf1
@@ -49,7 +49,7 @@
     @test n_ter.children[1] === n_ter_leaf1
     @test n_ter.children[2] === n_ter_leaf2
     @test n_ter.children[3] === n_ter_leaf3
-    @test DynamicExpressions.NodeModule.children(n_ter, Val(3)) ==
+    @test DynamicExpressions.NodeModule.get_children(n_ter, Val(3)) ==
         (n_ter_leaf1, n_ter_leaf2, n_ter_leaf3)
     @test n_ter.l === n_ter_leaf1
     @test n_ter.r === n_ter_leaf2
@@ -461,7 +461,7 @@ end
     @test readonly_tree.degree == 3
     @test readonly_tree.op == 1
 
-    ro_children = DynamicExpressions.NodeModule.children(readonly_tree, Val(3))
+    ro_children = DynamicExpressions.NodeModule.get_children(readonly_tree, Val(3))
     @test length(ro_children) == 3
     @test ro_children[1] isa DynamicExpressions.ReadOnlyNodeModule.AbstractReadOnlyNode
     @test ro_children[1].feature == 1