feat: guard undefined children behind Nullable

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -28,6 +28,7 @@ end
 import Reexport: @reexport
 macro ignore(args...) end
 
+import .UtilsModule: Nullable
 import .ValueInterfaceModule:
     is_valid,
     is_valid_array,

src/Interfaces.jl

@@ -3,6 +3,7 @@ module InterfacesModule
 
 using Interfaces: Interfaces, @interface, @implements, Arguments
 using DispatchDoctor: @unstable
+using ..UtilsModule: Nullable
 using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
 using ..NodeModule:
     Node,
@@ -14,6 +15,7 @@ using ..NodeModule:
     with_type_parameters,
     with_max_degree,
     max_degree,
+    unsafe_get_children,
     get_children,
     leaf_copy,
     leaf_convert,
@@ -228,7 +230,7 @@ function _check_create_node(tree::AbstractExpressionNode)
 end
 function _check_get_children(tree::AbstractExpressionNode{T,D}) where {T,D}
     tree.degree == 0 && return true
-    return get_children(tree) isa Tuple{typeof(tree),Vararg{typeof(tree)}} &&
+    return unsafe_get_children(tree) isa NTuple{D,Nullable{typeof(tree)}} &&
            get_children(tree, Val(D)) isa Tuple &&
            length(get_children(tree, Val(D))) == D &&
            length(get_children(tree, Val(1))) == 1

src/Node.jl

@@ -2,7 +2,7 @@ module NodeModule
 
 using DispatchDoctor: @unstable
 
-import ..UtilsModule: deprecate_varmap, Undefined
+import ..UtilsModule: deprecate_varmap, Undefined, Nullable
 
 const DEFAULT_NODE_TYPE = Float32
 const DEFAULT_MAX_DEGREE = 2
@@ -78,7 +78,7 @@ for N in (:Node, :GraphNode)
         val::T  # If is a constant, this stores the actual value
         feature::UInt16  # (Possibly undefined) If is a variable (e.g., x in cos(x)), this stores the feature index.
         op::UInt8  # (Possibly undefined) If operator, this is the index of the operator in the degree-specific operator enum
-        children::NTuple{D,$N{T,D}}
+        children::NTuple{D,Nullable{$N{T,D}}}
 
         #################
         ## Constructors:
@@ -173,13 +173,7 @@ Accessing this node should trigger some kind of noticable error
 (e.g., default returns itself, which causes infinite recursion).
 """
 function poison_node(n::AbstractNode)
-    # We don't want to use `nothing` because the type instability
-    # hits memory hard.
-    # Setting itself as the right child is the best thing,
-    # because it (1) doesn't allocate, and (2) will trigger
-    # infinite recursion errors if someone is mistakenly trying
-    # to access the right child when `.degree == 1`.
-    return n
+    return Nullable(true, n)
 end
 
 """
@@ -190,7 +184,7 @@ children may be "poisoned" nodes which you should not access,
 as they will trigger infinite recursion errors. Ensure to
 only access children only up to the `.degree` of the node.
 """
-@inline function get_children(node::AbstractNode)
+@inline function unsafe_get_children(node::AbstractNode)
     return getfield(node, :children)
 end
 
@@ -207,8 +201,8 @@ for total type stability.
     return get_children(node, Val(n))
 end
 @inline function get_children(node::AbstractNode{D}, ::Val{n}) where {D,n}
-    cs = get_children(node)
-    return ntuple(i -> cs[i], Val(n))
+    cs = unsafe_get_children(node)
+    return ntuple(i -> cs[i][], Val(n))
 end
 
 """
@@ -217,7 +211,7 @@ end
 Return the `i`-th child of a node (1-indexed).
 """
 @inline function get_child(n::AbstractNode{D}, i::Int) where {D}
-    return get_children(n)[i]
+    return unsafe_get_children(n)[i][]
 end
 
 """
@@ -227,7 +221,7 @@ Replace the `i`-th child of a node (1-indexed) with the given child node.
 Returns the new child. Updates the children tuple in-place.
 """
 @inline function set_child!(n::AbstractNode{D}, child::AbstractNode{D}, i::Int) where {D}
-    set_children!(n, Base.setindex(get_children(n), child, i))
+    set_children!(n, Base.setindex(unsafe_get_children(n), Nullable(false, child), i))
     return child
 end
 
@@ -242,17 +236,21 @@ provided than the node's maximum degree, remaining slots are filled with poison
 ) where {D}
     D2 = length(children)
     if D === D2
-        n.children = children
+        n.children = ntuple(i -> _ensure_nullable(@inbounds(children[i])), Val(D))
     else
         poison = poison_node(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 ? @inbounds(children[i]) : poison, Val(D))
+        # for simplicity, we can just use poison := Nullable(true, n)
+        # which will raise an UndefRefError if accessed.
+        n.children = ntuple(
+            i -> i <= D2 ? _ensure_nullable(@inbounds(children[i])) : poison, Val(D)
+        )
     end
 end
+@inline _ensure_nullable(x) = Nullable(false, x)
+@inline _ensure_nullable(x::Nullable) = x
 
 macro make_accessors(node_type)
     esc(
@@ -491,7 +489,7 @@ function set_leaf!(tree::AbstractExpressionNode, new_leaf::AbstractExpressionNod
 end
 function set_branch!(tree::AbstractExpressionNode, new_branch::AbstractExpressionNode)
     tree.op = new_branch.op
-    set_children!(tree, get_children(new_branch))
+    set_children!(tree, unsafe_get_children(new_branch))
     return nothing
 end
 

src/NodePreallocation.jl

@@ -56,7 +56,7 @@ function leaf_copy_into!(dest::N, src::N) where {N<:AbstractExpressionNode}
 end
 # COV_EXCL_STOP
 function branch_copy_into!(
-    dest::N, src::N, children::Vararg{N,M}
+    dest::N, src::N, children::Vararg{Any,M}
 ) where {T,D,N<:AbstractExpressionNode{T,D},M}
     dest.degree = M
     dest.op = src.op

src/NodeUtils.jl

@@ -1,5 +1,7 @@
 module NodeUtilsModule
 
+using ..UtilsModule: Nullable
+
 import ..NodeModule:
     AbstractNode,
     AbstractExpressionNode,
@@ -147,7 +149,7 @@ mutable struct NodeIndex{T,D} <: AbstractNode{D}
     degree::UInt8  # 0 for constant/variable, 1 for cos/sin, 2 for +/* etc.
     val::T  # If is a constant, this stores the actual value
     # ------------------- (possibly undefined below)
-    children::NTuple{D,NodeIndex{T,D}}
+    children::NTuple{D,Nullable{NodeIndex{T,D}}}
 
     function NodeIndex(::Type{_T}, ::Val{_D}, val) where {_T,_D}
         return new{_T,_D}(0, convert(_T, val))
@@ -166,10 +168,10 @@ NodeIndex(::Type{T}, ::Val{D}) where {T,D} = NodeIndex(T, Val(D), zero(T))
 
 @inline function Base.getproperty(n::NodeIndex, k::Symbol)
     if k == :l
-        # TODO: Should a depwarn be raised here? Or too slow?
-        return getfield(n, :children)[1]
+        # TODO: Should a deprecation warning be raised here? Or too slow?
+        return getfield(n, :children)[1][]
     elseif k == :r
-        return getfield(n, :children)[2]
+        return getfield(n, :children)[2][]
     else
         return getfield(n, k)
     end

src/ParametricExpression.jl

@@ -8,6 +8,7 @@ using ..NodeModule: AbstractExpressionNode, Node, tree_mapreduce
 using ..ExpressionModule:
     AbstractExpression, Metadata, with_contents, with_metadata, unpack_metadata
 using ..ChainRulesModule: NodeTangent
+using ..UtilsModule: Nullable
 
 import ..NodeModule:
     constructorof,
@@ -58,7 +59,7 @@ mutable struct ParametricNode{T,D} <: AbstractExpressionNode{T,D}
     parameter::UInt16  # Stores index of per-class parameter
 
     op::UInt8
-    children::NTuple{D,ParametricNode{T,D}}  # Children nodes
+    children::NTuple{D,Nullable{ParametricNode{T,D}}}  # Children nodes
 
     function ParametricNode{_T,_D}() where {_T,_D}
         n = new{_T,_D}()

src/Utils.jl

@@ -58,4 +58,16 @@ struct ResultOk2{A<:AbstractArray,B<:AbstractArray}
     ok::Bool
 end
 
+struct Nullable{T}
+    null::Bool
+    x::T
+end
+@inline function Base.getindex(x::Nullable)
+    x.null && throw(UndefRefError())
+    return x.x
+end
+@inline function Base.convert(::Type{Nullable{T}}, x::Nullable) where {T}
+    return Nullable(x.null, convert(T, x.x))
+end
+
 end

test/test_custom_node_type.jl

@@ -1,22 +1,25 @@
 using DynamicExpressions
+using DynamicExpressions: Nullable
 using Test
 
 mutable struct MyCustomNode{A,B} <: AbstractNode{2}
     degree::Int
     val1::A
     val2::B
-    children::NTuple{2,MyCustomNode{A,B}}
+    children::NTuple{2,Nullable{MyCustomNode{A,B}}}
 
     MyCustomNode(val1, val2) = new{typeof(val1),typeof(val2)}(0, val1, val2)
     function MyCustomNode(val1, val2, l)
         n = MyCustomNode(val1, val2)
         poison = n
         n.degree = 1
-        n.children = (l, poison)
+        set_children!(n, (l, poison))
         return n
     end
     function MyCustomNode(val1, val2, l, r)
-        return new{typeof(val1),typeof(val2)}(2, val1, val2, (l, r))
+        n = new{typeof(val1),typeof(val2)}(2, val1, val2)
+        set_children!(n, (l, r))
+        return n
     end
 end
 
@@ -50,7 +53,7 @@ mutable struct MyCustomNode2{T} <: AbstractExpressionNode{T,2}
     val::T
     feature::UInt16
     op::UInt8
-    children::NTuple{2,Base.RefValue{MyCustomNode2{T}}}
+    children::NTuple{2,Nullable{MyCustomNode2{T}}}
 end
 
 @test_throws ErrorException MyCustomNode2()

test/test_extra_node_fields.jl

@@ -2,7 +2,7 @@
 
 using Test
 using DynamicExpressions
-using DynamicExpressions: constructorof, max_degree
+using DynamicExpressions: Nullable, constructorof, max_degree
 
 mutable struct FrozenNode{T,D} <: AbstractExpressionNode{T,D}
     degree::UInt8
@@ -11,7 +11,7 @@ mutable struct FrozenNode{T,D} <: AbstractExpressionNode{T,D}
     frozen::Bool  # Extra field!
     feature::UInt16
     op::UInt8
-    children::NTuple{D,FrozenNode{T,D}}
+    children::NTuple{D,Nullable{FrozenNode{T,D}}}
 
     function FrozenNode{_T,_D}() where {_T,_D}
         n = new{_T,_D}()
@@ -104,5 +104,5 @@ ex = parse_expression(
 
 @test string_tree(ex) == "x + sin(y + 2.1)"
 @test ex.tree.frozen == false
-@test ex.tree.children[2].frozen == true
-@test ex.tree.children[2].children[1].frozen == false
+@test ex.tree.children[2][].frozen == true
+@test ex.tree.children[2][].children[1][].frozen == false