feat: create n-arity operator enum

Author: MilesCranmer

src/Evaluate.jl

@@ -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},

src/OperatorEnum.jl

@@ -6,29 +6,49 @@ abstract type AbstractOperatorEnum end
     OperatorEnum
 
 Defines an enum over operators, along with their derivatives.
+
 # Fields
-- `binops`: A tuple of binary operators. Scalar input type.
-- `unaops`: A tuple of unary operators. Scalar input type.
+- `ops`: A tuple of operators, with index `i` corresponding to the operator tuple for a node of degree `i`.
 """
-struct OperatorEnum{B,U} <: AbstractOperatorEnum
-    binops::B
-    unaops::U
+struct OperatorEnum{OPS<:Tuple{Vararg{Tuple}}} <: AbstractOperatorEnum
+    ops::OPS
+end
+
+function OperatorEnum(binary_operators::Tuple, unary_operators::Tuple)
+    return OperatorEnum((unary_operators, binary_operators))
 end
 
 """
     GenericOperatorEnum
 
 Defines an enum over operators, along with their derivatives.
+
 # Fields
-- `binops`: A tuple of binary operators.
-- `unaops`: A tuple of unary operators.
+- `ops`: A tuple of operators, with index `i` corresponding to the operator tuple for a node of degree `i`.
 """
-struct GenericOperatorEnum{B,U} <: AbstractOperatorEnum
-    binops::B
-    unaops::U
+struct GenericOperatorEnum{OPS<:Tuple{Vararg{Tuple}}} <: AbstractOperatorEnum
+    ops::OPS
+end
+
+function GenericOperatorEnum(binops::Tuple, unaops::Tuple)
+    return GenericOperatorEnum((unaops, binops))
 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)
+    return getfield(op, :ops)[i]
+end
+@inline function Base.getproperty(op::AbstractOperatorEnum, k::Symbol)
+    if k == :unaops
+        return getfield(op, :ops)[1]
+    elseif k == :binops
+        ops = getfield(op, :ops)
+        return length(ops) > 1 ? ops[2] : ()
+    else
+        return getfield(op, k)
+    end
+end
+
 end