wip

Author: MilesCranmer

Project.toml

@@ -13,7 +13,6 @@ PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 
 [weakdeps]

src/Node.jl

@@ -1,23 +1,21 @@
 module NodeModule
 
 using DispatchDoctor: @unstable
-using StaticArrays: SizedVector
 
 import ..OperatorEnumModule: AbstractOperatorEnum
 import ..UtilsModule: @memoize_on, @with_memoize, deprecate_varmap, Undefined
 
 const DEFAULT_NODE_TYPE = Float32
 
 """
-    AbstractNode{D,shared}
+    AbstractNode{D}
 
-Abstract type for D-arity trees. If `shared`, the node type
-permits graph-like structures. Must have the following fields:
+Abstract type for D-arity trees. Must have the following fields:
 
 - `degree::Integer`: Degree of the node. Either 0, 1, or 2. If 1,
     then `l` needs to be defined as the left child. If 2,
     then `r` also needs to be defined as the right child.
-- `children`: A collection of D children nodes.
+- `children`: A collection of D references to children nodes.
 
 # Deprecated fields
 
@@ -29,7 +27,7 @@ permits graph-like structures. Must have the following fields:
 - `r::AbstractNode{D}`: Right child of the current node. Should only
     be defined if `degree == 2`.
 """
-abstract type AbstractNode{D,shared} end
+abstract type AbstractNode{D} end
 
 """
     AbstractExpressionNode{T,D} <: AbstractNode{D}
@@ -73,25 +71,27 @@ You likely do not need to, but you could choose to override the following:
 - `with_type_parameters`
 
 """
-abstract type AbstractExpressionNode{T,D,shared} <: AbstractNode{D,shared} end
-
-mutable struct GeneralNode{T,D,shared} <: AbstractExpressionNode{T,D,shared}
-    degree::UInt8  # 0 for constant/variable, 1 for cos/sin, 2 for +/* etc.
-    constant::Bool  # false if variable
-    val::T  # If is a constant, this stores the actual value
-    feature::UInt16  # If is a variable (e.g., x in cos(x)), this stores the feature index.
-    op::UInt8  # If operator, this is the index of the operator in the degree-specific operator enum
-    children::SizedVector{D,GeneralNode{T,D,shared}}  # Children nodes
-
-    #################
-    ## Constructors:
-    #################
-    GeneralNode{_T,_D,_shared}() where {_T,_D,_shared} = new{_T,_D,_shared}()
+abstract type AbstractExpressionNode{T,D} <: AbstractNode{D} end
+
+for N in (:Node, :GraphNode)
+    @eval mutable struct $N{T,D} <: AbstractExpressionNode{T,D}
+        degree::UInt8  # 0 for constant/variable, 1 for cos/sin, 2 for +/* etc.
+        constant::Bool  # false if variable
+        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,Base.RefValue{$N{T,D}}}  # Children nodes
+
+        #################
+        ## Constructors:
+        #################
+        $N{_T,_D}() where {_T,_D} = new{_T,_D}()
+    end
 end
 
 #! format: off
 """
-    Node{T} <: AbstractExpressionNode{T,2}
+    Node{T,D} <: AbstractExpressionNode{T,D}
 
 Node defines a symbolic expression stored in a binary tree.
 A single `Node` instance is one "node" of this tree, and
@@ -113,7 +113,7 @@ nodes, you can evaluate or print a given expression.
     operator in `operators.binops`. In other words, this is an enum
     of the operators, and is dependent on the specific `OperatorEnum`
     object. Only defined if `degree >= 1`
-- `children::SizedArray{D,Node{T,D}}`: Children of the node. Only defined up to `degree`
+- `children::NTuple{D,Base.RefValue{Node{T,D}}}`: Children of the node. Only defined up to `degree`
 
 # Constructors
 
@@ -130,13 +130,13 @@ You may also construct nodes via the convenience operators generated by creating
 You may also choose to specify a default memory allocator for the node other than simply `Node{T}()`
 in the `allocator` keyword argument.
 """
-const Node{T} = GeneralNode{T,2,false}
+Node
 
 
 """
-    GraphNode{T} <: AbstractExpressionNode{T}
+    GraphNode{T,D} <: AbstractExpressionNode{T,D}
 
-Exactly the same as [`Node{T}`](@ref), but with the assumption that some
+Exactly the same as [`Node{T,D}`](@ref), but with the assumption that some
 nodes will be shared. All copies of this graph-like structure will
 be performed with this assumption, to preserve structure of the graph.
 
@@ -164,7 +164,7 @@ This has the same constructors as [`Node{T}`](@ref). Shared nodes
 are created simply by using the same node in multiple places
 when constructing or setting properties.
 """
-const GraphNode{T} = GeneralNode{T,2,true}
+GraphNode
 
 ################################################################################
 #! format: on
@@ -177,30 +177,41 @@ function max_degree(::Type{N}) where {N<:AbstractExpressionNode}
 end
 
 @unstable constructorof(::Type{<:Node}) = Node
+@unstable constructorof(::Type{<:Node{T,D} where T}) where {D} = Node{T,D} where T
 @unstable constructorof(::Type{<:GraphNode}) = GraphNode
+@unstable constructorof(::Type{<:GraphNode{T,D} where T}) where {D} = GraphNode{T,D} where T
 
-with_type_parameters(::Type{<:Node}, ::Type{T}) where {T} = Node{T}
-with_type_parameters(::Type{<:GraphNode}, ::Type{T}) where {T} = GraphNode{T}
+with_type_parameters(::Type{<:Node}, ::Type{T}, ::Val{D}=Val(2)) where {T,D} = Node{T,D}
+with_type_parameters(::Type{<:GraphNode}, ::Type{T}, ::Val{D}=Val(2)) where {T,D} = GraphNode{T,D}
 
-default_allocator(::Type{<:Node}, ::Type{T}) where {T} = Node{T}()
-default_allocator(::Type{<:GraphNode}, ::Type{T}) where {T} = GraphNode{T}()
+default_allocator(::Type{<:Node}, ::Type{T}, ::Val{D}=Val(2)) where {T,D} = Node{T,D}()
+default_allocator(::Type{<:GraphNode}, ::Type{T}, ::Val{D}=Val(2)) where {T,D} = GraphNode{T,D}()
 
 """Trait declaring whether nodes share children or not."""
 preserve_sharing(::Union{Type{<:AbstractNode},AbstractNode}) = false
-function preserve_sharing(
-    ::Union{Type{<:G},G}
-) where {shared,G<:GeneralNode{T,D,shared} where {T,D}}
-    return shared
-end
+preserve_sharing(::Union{Type{<:GraphNode},GraphNode}) = true
 
 include("base.jl")
 
 #! format: off
 @inline function (::Type{N})(
-    ::Type{T1}=Undefined; val=nothing, feature=nothing, op=nothing, children=nothing, allocator::F=default_allocator,
+    ::Type{T1}=Undefined; kws...
 ) where {T1,N<:AbstractExpressionNode,F}
-    validate_not_all_defaults(N, val, feature, op, children)
-    return node_factory(N, T1, val, feature, op, children, allocator)
+end
+@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,D,N<:AbstractExpressionNode{T,D} where T,F}
+    _children = if l !== nothing && r === nothing
+        @assert children === nothing
+        (l,)
+    elseif l !== nothing && r !== nothing
+        @assert children === nothing
+        (l, r)
+    else
+        children
+    end
+    validate_not_all_defaults(N, val, feature, op, _children)
+    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}
     return nothing
@@ -209,7 +220,7 @@ function validate_not_all_defaults(::Type{N}, val, feature, op, children) where
     if val === nothing && feature === nothing && op === nothing && children === nothing
         error(
             "Encountered the call for $N() inside the generic constructor. "
-            * "Did you forget to define `$(Base.typename(N).wrapper){T}() where {T} = new{T}()`?"
+            * "Did you forget to define `$(Base.typename(N).wrapper){T,D}() where {T,D} = new{T,D}()`?"
         )
     end
     return nothing
@@ -219,7 +230,7 @@ end
     ::Type{N}, ::Type{T1}, val::T2, ::Nothing, ::Nothing, ::Nothing, allocator::F,
 ) where {N,T1,T2,F}
     T = node_factory_type(N, T1, T2)
-    n = allocator(N, T)
+    n = allocator(N, T, D)
     n.degree = 0
     n.constant = true
     n.val = convert(T, val)
@@ -230,7 +241,7 @@ end
     ::Type{N}, ::Type{T1}, ::Nothing, feature::Integer, ::Nothing, ::Nothing, allocator::F,
 ) where {N,T1,F}
     T = node_factory_type(N, T1, DEFAULT_NODE_TYPE)
-    n = allocator(N, T)
+    n = allocator(N, T, D)
     n.degree = 0
     n.constant = false
     n.feature = feature
@@ -239,19 +250,14 @@ end
 """Create an operator node."""
 @inline function node_factory(
     ::Type{N}, ::Type, ::Nothing, ::Nothing, op::Integer, children::NTuple{D2}, allocator::F,
-) where {N,F,D2}
-    D = max_degree(N)
-    @assert D2 <= D
+) where {D,N<:AbstractExpressionNode{T where T,D},F,D2}
     T = promote_type(map(eltype, children)...)  # Always prefer existing nodes, so we don't mess up references from conversion
-    NT = with_type_parameters(N, T)
-    n = allocator(N, T)
+    NT = with_type_parameters(N, T, D)
+    n = allocator(N, T, D)
     n.degree = D2
     n.op = op
-    ar = SizedVector{D,NT}(undef)
-    for i in 1:D2
-        ar[i] = children[i]
-    end
-    n.children = ar
+    n.children
+    # map(Ref, children)
     return n
 end