feat: set_children! to work with vector

Author: MilesCranmer

src/Node.jl

@@ -199,7 +199,8 @@ end
     set_children!(n, Base.setindex(get_children(n), child, i))
     return child
 end
-@inline function set_children!(n::AbstractNode{D}, children::Tuple{Vararg{AbstractNode{D},D2}}) where {D,D2}
+@inline function set_children!(n::AbstractNode{D}, children::Union{Tuple,AbstractVector{<:AbstractNode{D}}}) where {D}
+    D2 = length(children)
     if D === D2
         n.children = children
     else
@@ -209,7 +210,7 @@ end
         # 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))
+        n.children = ntuple(i -> i <= D2 ? @inbounds(children[i]) : poison, Val(D))
     end
 end
 

src/ReadOnlyNode.jl

@@ -8,14 +8,7 @@ import ..NodeModule: default_allocator, with_type_parameters, constructorof, get
 abstract type AbstractReadOnlyNode{T,D,N<:AbstractExpressionNode{T,D}} <:
               AbstractExpressionNode{T,D} end
 
-"""A type of expression node that prevents writing to the inner node"""
-struct ReadOnlyNode{T,D,N} <: AbstractReadOnlyNode{T,D,N}
-    _inner::N
-
-    ReadOnlyNode(n::N) where {T,N<:AbstractExpressionNode{T}} = new{T,max_degree(N),N}(n)
-end
 @inline inner(n::AbstractReadOnlyNode) = getfield(n, :_inner)
-@unstable constructorof(::Type{<:ReadOnlyNode}) = ReadOnlyNode
 @inline function Base.getproperty(n::AbstractReadOnlyNode, s::Symbol)
     out = getproperty(inner(n), s)
     if out isa AbstractExpressionNode
@@ -25,12 +18,20 @@ end
     end
 end
 @inline function get_children(node::AbstractReadOnlyNode)
-    return map(ReadOnlyNode, get_children(inner(node)))
+    return map(constructorof(typeof(node)), get_children(inner(node)))
 end
 function Base.setproperty!(::AbstractReadOnlyNode, ::Symbol, v)
     return error("Cannot set properties on a ReadOnlyNode")
 end
 Base.propertynames(n::AbstractReadOnlyNode) = propertynames(inner(n))
-Base.copy(n::AbstractReadOnlyNode) = ReadOnlyNode(copy(inner(n)))
+Base.copy(n::AbstractReadOnlyNode) = constructorof(typeof(n))(copy(inner(n)))
+
+"""A type of expression node that prevents writing to the inner node"""
+struct ReadOnlyNode{T,D,N} <: AbstractReadOnlyNode{T,D,N}
+    _inner::N
+
+    ReadOnlyNode(n::N) where {T,N<:AbstractExpressionNode{T}} = new{T,max_degree(N),N}(n)
+end
+@unstable constructorof(::Type{<:ReadOnlyNode}) = ReadOnlyNode
 
 end

test/runtests.jl

@@ -18,7 +18,9 @@ end
 if "jet" in test_name
     @safetestset "JET" begin
         using Preferences
-        set_preferences!("DynamicExpressions", "dispatch_doctor_mode" => "disable"; force=true)
+        set_preferences!(
+            "DynamicExpressions", "dispatch_doctor_mode" => "disable"; force=true
+        )
         using JET
         using DynamicExpressions
         struct MyIgnoredModule