feat: n-arity strings

Author: MilesCranmer

src/Strings.jl

@@ -2,28 +2,43 @@ module StringsModule
 
 using ..UtilsModule: deprecate_varmap
 using ..OperatorEnumModule: AbstractOperatorEnum
-using ..NodeModule: AbstractExpressionNode, tree_mapreduce
+using ..NodeModule: AbstractExpressionNode, tree_mapreduce, max_degree
 
 function dispatch_op_name(
     ::Val{deg}, ::Nothing, idx, pretty::Bool
 )::Vector{Char} where {deg}
-    return vcat(
-        collect(deg == 1 ? "unary_operator[" : "binary_operator["),
-        collect(string(idx)),
-        [']'],
-    )
+    return vcat(collect(
+        if deg == 1
+            "unary_operator["
+        elseif deg == 2
+            "binary_operator["
+        else
+            "operator_deg$deg["
+        end,
+    ), collect(string(idx)), [']'])
 end
 function dispatch_op_name(
     ::Val{deg}, operators::AbstractOperatorEnum, idx, pretty::Bool
 ) where {deg}
-    op = if deg == 1
-        operators.unaops[idx]
-    else
-        operators.binops[idx]
-    end
+    op = operators[deg][idx]
     return collect((pretty ? get_pretty_op_name(op) : get_op_name(op))::String)
 end
 
+struct OpNameDispatcher{D,O<:AbstractOperatorEnum} <: Function
+    operators::O
+    pretty::Bool
+end
+@generated function (f::OpNameDispatcher{D,O})(branch) where {D,O}
+    return quote
+        degree = branch.degree
+        Base.Cartesian.@nif(
+            $D,
+            d -> d == degree,
+            d -> dispatch_op_name(Val(d), f.operators, branch.op, f.pretty),
+        )::Vector{Char}
+    end
+end
+
 const OP_NAME_CACHE = (; x=Dict{UInt64,String}(), lock=Threads.SpinLock())
 
 function get_op_name(op::F) where {F}
@@ -89,36 +104,30 @@ function string_variable(feature, variable_names)
 end
 
 # Vector of chars is faster than strings, so we use that.
-function combine_op_with_inputs(op, l, r)::Vector{Char}
-    if first(op) in ('+', '-', '*', '/', '^', '.', '>', '<', '=') || op == "!="
+function combine_op_with_inputs(op, args::Vararg{Any,D})::Vector{Char} where {D}
+    if D == 2 && (first(op) in ('+', '-', '*', '/', '^', '.', '>', '<', '=') || op == "!=")
         # "(l op r)"
         out = ['(']
-        append!(out, l)
+        append!(out, args[1])
         push!(out, ' ')
         append!(out, op)
         push!(out, ' ')
-        append!(out, r)
+        append!(out, args[2])
         push!(out, ')')
     else
         # "op(l, r)"
         out = copy(op)
         push!(out, '(')
-        append!(out, strip_brackets(l))
-        push!(out, ',')
-        push!(out, ' ')
-        append!(out, strip_brackets(r))
+        for i in 1:(D - 1)
+            append!(out, strip_brackets(args[i]))
+            push!(out, ',')
+            push!(out, ' ')
+        end
+        append!(out, strip_brackets(args[D]))
         push!(out, ')')
         return out
     end
 end
-function combine_op_with_inputs(op, l)
-    # "op(l)"
-    out = copy(op)
-    push!(out, '(')
-    append!(out, strip_brackets(l))
-    push!(out, ')')
-    return out
-end
 
 """
     string_tree(
@@ -169,13 +178,7 @@ function string_tree(
                 collect(f_variable(leaf.feature, variable_names))::Vector{Char}
             end
         end,
-        let operators = operators
-            (branch,) -> if branch.degree == 1
-                dispatch_op_name(Val(1), operators, branch.op, pretty)::Vector{Char}
-            else
-                dispatch_op_name(Val(2), operators, branch.op, pretty)::Vector{Char}
-            end
-        end,
+        OpNameDispatcher{max_degree(tree),typeof(operators)}(operators, pretty),
         combine_op_with_inputs,
         tree,
         Vector{Char};