feat: get expression algebra working

Author: MilesCranmer

src/Expression.jl

@@ -8,7 +8,8 @@ using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
 using ..UtilsModule: Undefined
 using ..ChainRulesModule: NodeTangent
 
-import ..NodeModule: copy_node, set_node!, count_nodes, tree_mapreduce, constructorof
+import ..NodeModule:
+    copy_node, set_node!, count_nodes, tree_mapreduce, constructorof, max_degree
 import ..NodeUtilsModule:
     preserve_sharing,
     count_constant_nodes,
@@ -99,9 +100,14 @@ end
     return Expression(tree, Metadata(d))
 end
 
+has_node_type(::Union{E,Type{E}}) where {N,E<:AbstractExpression{<:Any,N}} = true
+has_node_type(::Union{E,Type{E}}) where {E<:AbstractExpression} = false
 node_type(::Union{E,Type{E}}) where {N,E<:AbstractExpression{<:Any,N}} = N
+function max_degree(::Union{E,Type{E}}) where {E<:AbstractExpression}
+    return has_node_type(E) ? max_degree(node_type(E)) : max_degree(Node)
+end
 @unstable default_node_type(_) = Node
-default_node_type(::Type{<:AbstractExpression{T}}) where {T} = Node{T}
+default_node_type(::Type{N}) where {T,N<:AbstractExpression{T}} = Node{T,max_degree(N)}
 
 ########################################################
 # Abstract interface ###################################

src/ExpressionAlgebra.jl

@@ -54,19 +54,33 @@ of the expression.
 """
 declare_operator_alias(op::F, _) where {F<:Function} = op
 
+allow_chaining(@nospecialize(op)) = false
+allow_chaining(::typeof(+)) = true
+allow_chaining(::typeof(*)) = true
+
 function apply_operator(op::F, args::Vararg{Any,D}) where {F<:Function,D}
     idx = findfirst(e -> e isa AbstractExpression, args)::Int
     example_expr = args[idx]
     E = typeof(example_expr)
     @assert all(e -> !(e isa AbstractExpression) || typeof(e) === E, args)
     operators = get_operators(example_expr, nothing)
 
-    op_idx = findfirst(==(op), map(Base.Fix2(declare_operator_alias, Val(D)), operators[D]))
+    op_idx = if length(operators) >= D
+        findfirst(==(op), map(Base.Fix2(declare_operator_alias, Val(D)), operators[D]))
+    else
+        nothing
+    end
     if isnothing(op_idx)
+        if allow_chaining(op) && D > 2
+            # These operators might get chained by Julia, so we check
+            # downward for any matching arity.
+            inner = apply_operator(op, args[1:(end - 1)]...)
+            return apply_operator(op, inner, args[end])
+        end
         throw(
             MissingOperatorError(
                 "Operator $op not found in operators for expression type " *
-                "$(typeof(l)) with $(D)-degree operators $(operators[D])",
+                "$(E) with $(D)-degree operators $(operators[D])",
             ),
         )
     end
@@ -116,7 +130,6 @@ macro declare_expression_operator(op, arity)
             continue
         end
 
-
         arglist = [Expr(:(::), syms[i], types[i]) for i in 1:arity]
         signature = Expr(:call, op, arglist...)
         if any(t -> t == :T || (t isa Expr && t.head == :curly && :T in t.args), types)

src/OperatorEnum.jl

@@ -42,6 +42,9 @@ Base.copy(op::AbstractOperatorEnum) = op
 @unstable @inline function Base.getindex(op::AbstractOperatorEnum, i::Int)
     return getfield(op, :ops)[i]
 end
+@inline function Base.length(op::AbstractOperatorEnum)
+    return length(getfield(op, :ops))
+end
 @inline function Base.getproperty(op::AbstractOperatorEnum, k::Symbol)
     if k == :unaops
         return getfield(op, :ops)[1]

src/ParametricExpression.jl

@@ -122,7 +122,9 @@ function with_type_parameters(::Type{N}, ::Type{T}) where {N<:ParametricNode,T}
     return ParametricNode{T,max_degree(N)}
 end
 @unstable default_node_type(::Type{<:ParametricExpression}) = ParametricNode{T,2} where {T}
-default_node_type(::Type{<:ParametricExpression{T}}) where {T} = ParametricNode{T,2}
+function default_node_type(::Type{N}) where {T,N<:ParametricExpression{T}}
+    return ParametricNode{T,max_degree(N)}
+end
 preserve_sharing(::Union{Type{<:ParametricNode},ParametricNode}) = false # TODO: Change this?
 function leaf_copy(t::ParametricNode{T}) where {T}
     if t.constant

test/test_buffered_evaluation.jl

@@ -147,7 +147,7 @@ end
         result2, ok2 = eval_tree_array(tree, X, operators; eval_options)
 
         # Results should be identical
-        @test result1 ≈ result2
+        @test isapprox(result1, result2; atol=1e-10)
         @test ok1 == ok2
     end
 end