Permit n-argument operators

src/DynamicExpressions.jl

@@ -51,6 +51,7 @@ import .NodeModule:
     constructorof,
     with_type_parameters,
     preserve_sharing,
+    max_degree,
     leaf_copy,
     branch_copy,
     leaf_hash,

src/Evaluate.jl

@@ -2,7 +2,7 @@ module EvaluateModule
 
 using DispatchDoctor: @stable, @unstable
 
-import ..NodeModule: AbstractExpressionNode, constructorof
+import ..NodeModule: AbstractExpressionNode, constructorof, max_degree, children
 import ..StringsModule: string_tree
 import ..OperatorEnumModule: OperatorEnum, GenericOperatorEnum
 import ..UtilsModule: fill_similar, counttuple, ResultOk
@@ -251,8 +251,14 @@ end
 
 # These are marked unstable due to issues discussed on
 # https://github.com/JuliaLang/julia/issues/55147
-@unstable get_nuna(::Type{<:OperatorEnum{B,U}}) where {B,U} = counttuple(U)
-@unstable get_nbin(::Type{<:OperatorEnum{B}}) where {B} = counttuple(B)
+@unstable function get_nuna(::Type{<:OperatorEnum{OPS}}) where {OPS}
+    ts = OPS.types
+    return isempty(ts) ? 0 : counttuple(ts[1])
+end
+@unstable function get_nbin(::Type{<:OperatorEnum{OPS}}) where {OPS}
+    ts = OPS.types
+    return length(ts) == 1 ? 0 : counttuple(ts[2])
+end
 
 function _eval_tree_array(
     tree::AbstractExpressionNode{T},
@@ -275,35 +281,43 @@ function _eval_tree_array(
     elseif tree.degree == 1
         op_idx = tree.op
         return dispatch_deg1_eval(tree, cX, op_idx, operators, eval_options)
-    else
+    elseif max_degree(tree) == 2 || tree.degree == 2
         # TODO - add op(op2(x, y), z) and op(x, op2(y, z))
         # op(x, y), where x, y are constants or variables.
         op_idx = tree.op
         return dispatch_deg2_eval(tree, cX, op_idx, operators, eval_options)
+    else
+        op_idx = tree.op
+        return dispatch_degn_eval(tree, cX, op_idx, operators, eval_options)
+    end
+end
+
+@generated function degn_eval(
+    cumulators::NTuple{N,<:AbstractVector{T}}, op::F, ::EvalOptions{false}
+)::ResultOk where {N,T,F}
+    # Fast general implementation of `cumulators[1] .= op.(cumulators[1], cumulators[2], ...)`
+    quote
+        Base.Cartesian.@nexprs($N, i -> cumulator_i = cumulators[i])
+        @inbounds @simd for j in eachindex(cumulator_1)
+            cumulator_1[j] = Base.Cartesian.@ncall($N, op, i -> cumulator_i[j])::T
+        end
+        return ResultOk(cumulator_1, true)
     end
 end
 
 function deg2_eval(
     cumulator_l::AbstractVector{T},
     cumulator_r::AbstractVector{T},
     op::F,
-    ::EvalOptions{false},
+    eval_options::EvalOptions{false},
 )::ResultOk where {T,F}
-    @inbounds @simd for j in eachindex(cumulator_l)
-        x = op(cumulator_l[j], cumulator_r[j])::T
-        cumulator_l[j] = x
-    end
-    return ResultOk(cumulator_l, true)
+    return degn_eval((cumulator_l, cumulator_r), op, eval_options)
 end
 
 function deg1_eval(
-    cumulator::AbstractVector{T}, op::F, ::EvalOptions{false}
+    cumulator::AbstractVector{T}, op::F, eval_options::EvalOptions{false}
 )::ResultOk where {T,F}
-    @inbounds @simd for j in eachindex(cumulator)
-        x = op(cumulator[j])::T
-        cumulator[j] = x
-    end
-    return ResultOk(cumulator, true)
+    return degn_eval((cumulator,), op, eval_options)
 end
 
 function deg0_eval(
@@ -318,6 +332,52 @@ function deg0_eval(
     end
 end
 
+@generated function inner_dispatch_degn_eval(
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    op_idx::Integer,
+    ::Val{degree},
+    operators::OperatorEnum{OPS},
+    eval_options::EvalOptions,
+) where {T,degree,OPS}
+    nops = length(OPS.types[degree].types)
+    return quote
+        cs = children(tree, Val($degree))
+        Base.Cartesian.@nexprs(
+            $degree,
+            i -> begin
+                result_i = _eval_tree_array(cs[i], cX, operators, eval_options)
+                !result_i.ok && return result_i
+                @return_on_nonfinite_array(eval_options, result_i.x)
+            end
+        )
+        cumulators = Base.Cartesian.@ntuple($degree, i -> result_i.x)
+        Base.Cartesian.@nif(
+            $nops,
+            i -> i == op_idx,
+            i -> degn_eval(cumulators, operators[$degree][i], eval_options),
+        )
+    end
+end
+@generated function dispatch_degn_eval(
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    op_idx::Integer,
+    operators::OperatorEnum,
+    eval_options::EvalOptions,
+) where {T}
+    D = max_degree(tree)
+    return quote
+        # If statement over degrees
+        degree = tree.degree
+        return Base.Cartesian.@nif(
+            $D,
+            d -> d == degree,
+            d ->
+                inner_dispatch_degn_eval(tree, cX, op_idx, Val(d), operators, eval_options)
+        )
+    end
+end
 @generated function dispatch_deg2_eval(
     tree::AbstractExpressionNode{T},
     cX::AbstractMatrix{T},

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,86 +54,96 @@ of the expression.
 """
 declare_operator_alias(op::F, _) where {F<:Function} = op
 
-function apply_operator(op::F, l::AbstractExpression) where {F<:Function}
-    operators = get_operators(l, nothing)
-    op_idx = findfirst(
-        ==(op), map(Base.Fix2(declare_operator_alias, Val(1)), operators.unaops)
-    )
-    if op_idx === nothing
-        throw(
-            MissingOperatorError(
-                "Operator $op not found in operators for expression type $(typeof(l)) with unary operators $(operators.unaops)",
-            ),
-        )
-    end
-    return insert_operator_index(op_idx, (l,), l)
-end
-function apply_operator(op::F, l, r) where {F<:Function}
-    (operators, example_expr) = if l isa AbstractExpression && r isa AbstractExpression
-        @assert typeof(r) === typeof(l)
-        (get_operators(l, nothing), l)
-    elseif l isa AbstractExpression
-        (get_operators(l, nothing), l)
+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 = if length(operators) >= D
+        findfirst(==(op), map(Base.Fix2(declare_operator_alias, Val(D)), operators[D]))
     else
-        r::AbstractExpression
-        (get_operators(r, nothing), r)
+        nothing
     end
-    op_idx = findfirst(
-        ==(op), map(Base.Fix2(declare_operator_alias, Val(2)), operators.binops)
-    )
-    if op_idx === nothing
+    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 binary operators $(operators.binops)",
+                "Operator $op not found in operators for expression type " *
+                "$(E) with $(D)-degree operators $(operators[D])",
             ),
         )
     end
-    return insert_operator_index(op_idx, (l, r), example_expr)
+    return insert_operator_index(op_idx, args, example_expr)
 end
 
 """
     @declare_expression_operator(op, arity)
 
 Declare an operator function for `AbstractExpression` types.
 
-This macro generates a method for the given operator `op` that works with
-`AbstractExpression` arguments. The `arity` parameter specifies whether
-the operator is unary (1) or binary (2).
-
-# Arguments
-- `op`: The operator to be declared (e.g., `Base.sin`, `Base.:+`).
-- `arity`: The number of arguments the operator takes (1 for unary, 2 for binary).
+This macro generates methods for the given operator `op` that work with
+`AbstractExpression` arguments.  The `arity` parameter specifies the number
+of arguments the operator takes.
 """
 macro declare_expression_operator(op, arity)
-    @assert arity ∈ (1, 2)
+    syms = [Symbol('x', i) for i in 1:arity]
+    AE = :($(AbstractExpression))
     if arity == 1
         return esc(
             quote
-                $op(l::AbstractExpression) = $(apply_operator)($op, l)
+                $op($(only(syms))::$(AE)) = $(apply_operator)($op, $(only(syms)))
             end,
         )
-    elseif arity == 2
-        return esc(
-            quote
-                function $op(l::AbstractExpression, r::AbstractExpression)
-                    return $(apply_operator)($op, l, r)
-                end
-                function $op(l::T, r::AbstractExpression{T}) where {T}
-                    return $(apply_operator)($op, l, r)
-                end
-                function $op(l::AbstractExpression{T}, r::T) where {T}
-                    return $(apply_operator)($op, l, r)
-                end
-                # Convenience methods for Number types
-                function $op(l::Number, r::AbstractExpression{T}) where {T}
-                    return $(apply_operator)($op, l, r)
-                end
-                function $op(l::AbstractExpression{T}, r::Number) where {T}
-                    return $(apply_operator)($op, l, r)
-                end
-            end,
+    end
+
+    wrappers = (AE, :($(AE){T}), :T, :Number)
+    methods = Expr(:block)
+
+    for types in Iterators.product(ntuple(_ -> wrappers, arity)...)
+        has_expr = any(
+            t -> t == AE || (t isa Expr && t.head == :curly && t.args[1] == AE), types
+        )
+        has_plain_T = any(==(:T), types)
+        has_abstract_expr_T = any(
+            t -> t isa Expr && t.head == :curly && t.args[1] == AE && :T in t.args, types
         )
+        has_abstract_expr_plain = any(==(AE), types)
+        if any((
+            !has_expr,
+            # ^At least one arg must be an AbstractExpression (avoid type‑piracy)
+            has_abstract_expr_plain && has_abstract_expr_T,
+            # ^If a plain `T` appears, ensure an `AbstractExpression{T}` is also present
+            has_plain_T ⊻ has_abstract_expr_T,
+            # ^Do not mix bare `AbstractExpression` with `AbstractExpression{T}`
+        ))
+            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)
+            signature = Expr(:where, signature, :(T))
+        end
+
+        body = Expr(:block, :(return $(apply_operator)($op, $(syms...))))
+
+        fn = Expr(:function, signature, body)
+
+        push!(methods.args, fn)
     end
+
+    return esc(methods)
 end
 
 #! format: off
@@ -159,6 +169,11 @@ for op in (
 )
     @eval @declare_expression_operator Base.$(op) 2
 end
+for op in (
+    :*, :+, :clamp, :max, :min, :fma, :muladd,
+)
+    @eval @declare_expression_operator Base.$(op) 3
+end
 #! format: on
 
 end

src/Interfaces.jl

@@ -70,8 +70,10 @@ function _check_get_metadata(ex::AbstractExpression)
     new_ex = with_metadata(ex, get_metadata(ex))
     return new_ex == ex && new_ex isa typeof(ex)
 end
-function _check_get_tree(ex::AbstractExpression{T,N}) where {T,N}
-    return get_tree(ex) isa N || get_tree(ex) isa AbstractReadOnlyNode{T,N}
+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}
 end
 function _check_get_operators(ex::AbstractExpression)
     return get_operators(ex) isa AbstractOperatorEnum
@@ -143,9 +145,11 @@ end
 function _check_constructorof(ex::AbstractExpression)
     return constructorof(typeof(ex)) isa Base.Callable
 end
-function _check_tree_mapreduce(ex::AbstractExpression{T,N}) where {T,N}
+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,N}}})
+           (Vector{N2} where {N2<:Union{N,AbstractReadOnlyNode{T,D,N}}})
 end
 
 #! format: off