feat: get expression operators working with 3-arg input

Author: MilesCranmer

src/ExpressionAlgebra.jl

@@ -54,86 +54,83 @@ 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
+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]))
+    if isnothing(op_idx)
         throw(
             MissingOperatorError(
-                "Operator $op not found in operators for expression type $(typeof(l)) with unary operators $(operators.unaops)",
+                "Operator $op not found in operators for expression type " *
+                "$(typeof(l)) with $(D)-degree operators $(operators[D])",
             ),
         )
     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)
-    else
-        r::AbstractExpression
-        (get_operators(r, nothing), r)
-    end
-    op_idx = findfirst(
-        ==(op), map(Base.Fix2(declare_operator_alias, Val(2)), operators.binops)
-    )
-    if op_idx === nothing
-        throw(
-            MissingOperatorError(
-                "Operator $op not found in operators for expression type $(typeof(l)) with binary operators $(operators.binops)",
-            ),
-        )
-    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 +156,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/OperatorEnum.jl

@@ -1,5 +1,7 @@
 module OperatorEnumModule
 
+using DispatchDoctor: @unstable
+
 abstract type AbstractOperatorEnum end
 
 """
@@ -37,7 +39,7 @@ end
 Base.copy(op::AbstractOperatorEnum) = op
 # TODO: Is this safe? What if a vector is passed here?
 
-@inline function Base.getindex(op::AbstractOperatorEnum, i::Int)
+@unstable @inline function Base.getindex(op::AbstractOperatorEnum, i::Int)
     return getfield(op, :ops)[i]
 end
 @inline function Base.getproperty(op::AbstractOperatorEnum, k::Symbol)

src/Strings.jl

@@ -24,7 +24,7 @@ function dispatch_op_name(
     return collect((pretty ? get_pretty_op_name(op) : get_op_name(op))::String)
 end
 
-struct OpNameDispatcher{D,O<:AbstractOperatorEnum} <: Function
+struct OpNameDispatcher{D,O<:Union{AbstractOperatorEnum,Nothing}} <: Function
     operators::O
     pretty::Bool
 end