refactor: remove ReadOnlyNode and simplify interface

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -22,7 +22,6 @@ using DispatchDoctor: @stable, @unstable
     include("Random.jl")
     include("Parse.jl")
     include("ParametricExpression.jl")
-    include("ReadOnlyNode.jl")
     include("StructuredExpression.jl")
 end
 
@@ -103,7 +102,6 @@ import .ExpressionAlgebraModule: declare_operator_alias
 @reexport import .ParseModule: @parse_expression, parse_expression
 import .ParseModule: parse_leaf
 @reexport import .ParametricExpressionModule: ParametricExpression, ParametricNode
-import .ReadOnlyNodeModule: ReadOnlyNode
 @reexport import .StructuredExpressionModule: StructuredExpression
 import .StructuredExpressionModule: AbstractStructuredExpression
 

src/Interfaces.jl

@@ -57,7 +57,6 @@ using ..ExpressionModule:
     with_metadata,
     default_node_type
 using ..ParametricExpressionModule: ParametricExpression, ParametricNode
-using ..ReadOnlyNodeModule: AbstractReadOnlyNode
 using ..StructuredExpressionModule: StructuredExpression
 
 ###############################################################################
@@ -76,7 +75,7 @@ end
 function _check_get_tree(
     ex::AbstractExpression{T,N}
 ) where {T,D,N<:AbstractExpressionNode{T,D}}
-    return get_tree(ex) isa N || get_tree(ex) isa AbstractReadOnlyNode{T,D,N}
+    return get_tree(ex) isa N
 end
 function _check_get_operators(ex::AbstractExpression)
     return get_operators(ex) isa AbstractOperatorEnum
@@ -151,8 +150,7 @@ end
 function _check_tree_mapreduce(
     ex::AbstractExpression{T,N}
 ) where {T,D,N<:AbstractExpressionNode{T,D}}
-    return tree_mapreduce(node -> [node], vcat, ex) isa
-           (Vector{N2} where {N2<:Union{N,AbstractReadOnlyNode{T,D,N}}})
+    return tree_mapreduce(node -> [node], vcat, ex) isa Vector{<:N}
 end
 
 #! format: off

src/Node.jl

@@ -298,25 +298,30 @@ Base.eltype(::AbstractExpressionNode{T}) where {T} = T
 
 has_max_degree(::Type{<:AbstractNode}) = false
 has_max_degree(::Type{<:AbstractNode{D}}) where {D} = true
+has_eltype(::Type{<:AbstractExpressionNode}) = false
+has_eltype(::Type{<:AbstractExpressionNode{T}}) where {T} = true
 # COV_EXCL_STOP
 #! format: on
 
-@unstable function constructorof(::Type{N}) where {N<:Node}
-    return Node{T,max_degree(N)} where {T}
+@unstable function node_wrapper(::Type{N}) where {N<:AbstractExpressionNode}
+    return Base.typename(N).wrapper
 end
-@unstable function constructorof(::Type{N}) where {N<:GraphNode}
-    return GraphNode{T,max_degree(N)} where {T}
+@unstable function constructorof(::Type{N}) where {N<:AbstractExpressionNode}
+    return node_wrapper(N){T,max_degree(N)} where {T}
 end
-
-function with_type_parameters(::Type{N}, ::Type{T}) where {N<:Node,T}
-    return Node{T,max_degree(N)}
+function with_type_parameters(::Type{N}, ::Type{T}) where {N<:AbstractExpressionNode,T}
+    return node_wrapper(N){T,max_degree(N)}
 end
-function with_type_parameters(::Type{N}, ::Type{T}) where {N<:GraphNode,T}
-    return GraphNode{T,max_degree(N)}
+@unstable function with_max_degree(::Type{N}, ::Val{D}) where {N<:AbstractExpressionNode,D}
+    if has_eltype(N)
+        return node_wrapper(N){eltype(N),D}
+    else
+        return node_wrapper(N){T,D} where {T}
+    end
+end
+@unstable function with_default_max_degree(::Type{N}) where {N<:AbstractNode}
+    return with_max_degree(N, Val(max_degree(N)))
 end
-
-with_max_degree(::Type{N}, ::Val{D}) where {T,N<:Node{T},D} = Node{T,D}
-with_max_degree(::Type{N}, ::Val{D}) where {T,N<:GraphNode{T},D} = GraphNode{T,D}
 
 function default_allocator(::Type{N}, ::Type{T}) where {N<:AbstractExpressionNode,T}
     return with_type_parameters(N, T)()
@@ -341,17 +346,26 @@ include("base.jl")
     else
         children
     end
-    validate_not_all_defaults(N, val, feature, op, _children)
+    if all_defaults(N, val, feature, op, _children)
+        return make_default(N, T1)
+    end
     return node_factory(N, T1, val, feature, op, _children, allocator)
 end
-function validate_not_all_defaults(::Type{N}, val, feature, op, children) where {N<:AbstractExpressionNode}
-    if all(isnothing, (val, feature, op, children))
+function make_default(::Type{N}, ::Type{T1}) where {T1,N<:AbstractExpressionNode}
+    if has_max_degree(N)
         error(
             "Encountered the call for $N() inside the generic constructor. "
             * "Did you forget to define `$(Base.typename(N).wrapper){T,D}() where {T,D} = new{T,D}()`?"
         )
     end
-    return nothing
+    if T1 === Undefined
+        return with_default_max_degree(N)()
+    else
+        return with_type_parameters(with_default_max_degree(N), T1)()
+    end
+end
+function all_defaults(::Type{N}, val, feature, op, children) where {N<:AbstractExpressionNode}
+    return all(isnothing, (val, feature, op, children))
 end
 """Create a constant leaf."""
 @inline function node_factory(

src/ParametricExpression.jl

@@ -13,6 +13,7 @@ import ..NodeModule:
     constructorof,
     with_type_parameters,
     with_max_degree,
+    with_default_max_degree,
     max_degree,
     preserve_sharing,
     get_children,
@@ -118,16 +119,10 @@ end
 ###############################################################################
 # Abstract expression node interface ##########################################
 ###############################################################################
-@unstable constructorof(::Type{N}) where {N<:ParametricNode} =
-    ParametricNode{T,max_degree(N)} where {T}
 @unstable constructorof(::Type{<:ParametricExpression}) = ParametricExpression
-function with_type_parameters(::Type{N}, ::Type{T}) where {N<:ParametricNode,T}
-    return ParametricNode{T,max_degree(N)}
-end
-function with_max_degree(::Type{N}, ::Val{D}) where {T,N<:ParametricNode{T},D}
-    return ParametricNode{T,D}
+@unstable function default_node_type(::Type{<:ParametricExpression})
+    return with_default_max_degree(ParametricNode)
 end
-@unstable default_node_type(::Type{<:ParametricExpression}) = ParametricNode{T,2} where {T}
 function default_node_type(::Type{N}) where {T,N<:ParametricExpression{T}}
     return ParametricNode{T,max_degree(N)}
 end

src/ReadOnlyNode.jl

@@ -21,7 +21,6 @@ import ..ExpressionModule:
     node_type,
     get_scalar_constants,
     set_scalar_constants!
-import ..ReadOnlyNodeModule: ReadOnlyNode
 
 abstract type AbstractStructuredExpression{
     T,F<:Function,N<:AbstractExpressionNode{T},E<:AbstractExpression{T,N},D<:NamedTuple
@@ -132,7 +131,7 @@ function get_metadata(e::AbstractStructuredExpression)
     return e.metadata
 end
 function get_tree(e::AbstractStructuredExpression)
-    return ReadOnlyNode(get_tree(get_metadata(e).structure(get_contents(e))))
+    return get_tree(get_metadata(e).structure(get_contents(e)))
 end
 function get_operators(
     e::AbstractStructuredExpression, operators::Union{AbstractOperatorEnum,Nothing}=nothing

src/StructuredExpression.jl

@@ -448,44 +448,6 @@ end
     @test node_from_pex.children[3].feature == 3
 end
 
-@testitem "ReadOnlyNode with N-ary Node" tags = [:narity] begin
-    using DynamicExpressions
-    using Test
-
-    my_ro_unary_op(x) = x
-    my_ro_binary_op(x, y) = x
-    my_ro_ternary_op(x, y, z) = x
-
-    operators_ro = OperatorEnum(
-        1 => (my_ro_unary_op,), 2 => (my_ro_binary_op,), 3 => (my_ro_ternary_op,)
-    )
-    DynamicExpressions.@extend_operators operators_ro
-
-    x1_ro = Node{Float64,3}(; feature=1)
-    x2_ro = Node{Float64,3}(; feature=2)
-    x3_ro = Node{Float64,3}(; feature=3)
-    tree_ro_ter = Node{Float64,3}(; op=1, children=(x1_ro, x2_ro, x3_ro))
-
-    expr_ro = Expression(tree_ro_ter; operators=operators_ro)
-    readonly_tree = DynamicExpressions.ReadOnlyNode(DynamicExpressions.get_tree(expr_ro))
-
-    @test readonly_tree isa DynamicExpressions.ReadOnlyNodeModule.AbstractReadOnlyNode
-    inner_node_ro = DynamicExpressions.ReadOnlyNodeModule.inner(readonly_tree)
-    @test DynamicExpressions.NodeModule.max_degree(inner_node_ro) == 3
-    @test readonly_tree.degree == 3
-    @test readonly_tree.op == 1
-
-    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
-    @test ro_children[2].feature == 2
-    @test ro_children[3].feature == 3
-
-    @test readonly_tree.l.feature == 1
-    @test readonly_tree.r.feature == 2
-end
-
 @testitem "NodeUtils.jl NodeIndex for N-ary" tags = [:narity] begin
     using DynamicExpressions
     using Test

test/test_n_arity_nodes.jl

@@ -15,9 +15,7 @@ end
 
 @testitem "OperatorEnumConstruction helper internals & edge cases" begin
     using DynamicExpressions.OperatorEnumConstructionModule:
-        _unpack_broadcast_function,
-        OperatorEnum,
-        empty_all_globals!
+        _unpack_broadcast_function, OperatorEnum, empty_all_globals!
     using Base.Broadcast: BroadcastFunction
     using Test
 

test/test_operator_construction_edgecases.jl

@@ -131,7 +131,6 @@ include("test_operator_construction_edgecases.jl")
 include("test_node_interface.jl")
 include("test_expression_math.jl")
 include("test_structured_expression.jl")
-include("test_readonlynode.jl")
 include("test_zygote_gradient_wrapper.jl")
 include("test_supposition_consistency.jl")
 include("test_n_arity_nodes.jl")