Merge pull request #97 from SymbolicML/prevent-julia-bugs

Increase purity of `@generated` functions during tests with DispatchDoctor

Author: MilesCranmer

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.19.0"
+version = "0.19.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/DynamicExpressions.jl

@@ -67,6 +67,7 @@ import .NodeModule:
     get_scalar_constants,
     set_scalar_constants!
 @reexport import .StringsModule: string_tree, print_tree
+import .StringsModule: get_op_name
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!

src/Evaluate.jl

@@ -52,11 +52,7 @@ struct EvalOptions{T,B,E}
     early_exit::Val{E}
 end
 
-@stable(
-    default_mode = "disable",
-    default_union_limit = 2,
-    @inline _to_bool_val(x::Bool) = x ? Val(true) : Val(false)
-)
+@unstable @inline _to_bool_val(x::Bool) = x ? Val(true) : Val(false)
 @inline _to_bool_val(x::Val{T}) where {T} = Val(T::Bool)
 
 @unstable function EvalOptions(;
@@ -74,11 +70,11 @@ end
         throw(ArgumentError("Invalid keyword argument(s): $(keys(deprecated_kws))"))
     end
     if !isempty(deprecated_kws)
-        @assert eval_options === nothing "Cannot use both `eval_options` and deprecated flags `turbo` and `bumper`."
-        Base.depwarn(
-            "The `turbo` and `bumper` keyword arguments are deprecated. Please use `eval_options` instead.",
-            :eval_tree_array,
+        @assert(
+            eval_options === nothing,
+            "Cannot use both `eval_options` and deprecated flags `turbo` and `bumper`."
         )
+        # TODO: We don't do a depwarn as it can GREATLY bottleneck the search speed.
     end
     if eval_options !== nothing
         return eval_options
@@ -183,8 +179,10 @@ function eval_tree_array(
     return eval_tree_array(tree, cX, operators; kws...)
 end
 
-get_nuna(::Type{<:OperatorEnum{B,U}}) where {B,U} = counttuple(U)
-get_nbin(::Type{<:OperatorEnum{B}}) where {B} = counttuple(B)
+# 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)
 
 function _eval_tree_array(
     tree::AbstractExpressionNode{T},

src/Strings.jl

@@ -4,16 +4,6 @@ using ..UtilsModule: deprecate_varmap
 using ..OperatorEnumModule: AbstractOperatorEnum
 using ..NodeModule: AbstractExpressionNode, tree_mapreduce
 
-const OP_NAMES = Base.ImmutableDict(
-    "safe_log" => "log",
-    "safe_log2" => "log2",
-    "safe_log10" => "log10",
-    "safe_log1p" => "log1p",
-    "safe_acosh" => "acosh",
-    "safe_sqrt" => "sqrt",
-    "safe_pow" => "^",
-)
-
 function dispatch_op_name(::Val{deg}, ::Nothing, idx)::Vector{Char} where {deg}
     if deg == 1
         return vcat(collect("unary_operator["), collect(string(idx)), [']'])
@@ -23,34 +13,38 @@ function dispatch_op_name(::Val{deg}, ::Nothing, idx)::Vector{Char} where {deg}
 end
 function dispatch_op_name(::Val{deg}, operators::AbstractOperatorEnum, idx) where {deg}
     if deg == 1
-        return get_op_name(operators.unaops[idx])::Vector{Char}
+        return collect(get_op_name(operators.unaops[idx])::String)
     else
-        return get_op_name(operators.binops[idx])::Vector{Char}
+        return collect(get_op_name(operators.binops[idx])::String)
     end
 end
 
-@generated function get_op_name(op::F)::Vector{Char} where {F}
+const OP_NAME_CACHE = (; x=Dict{UInt64,String}(), lock=Threads.SpinLock())
+
+function get_op_name(op::F) where {F}
+    h = hash(op)
+    lock(OP_NAME_CACHE.lock)
     try
-        # Bit faster to just cache the name of the operator:
-        op_s = if F <: Broadcast.BroadcastFunction
-            string(F.parameters[1].instance) * '.'
-        else
-            string(F.instance)
+        cache = OP_NAME_CACHE.x
+        if haskey(cache, h)
+            return cache[h]
         end
-        if length(op_s) == 2 && op_s[1] in ('+', '-', '*', '/', '^') && op_s[2] == '.'
-            op_s = '.' * op_s[1]
-        end
-        out = collect(get(OP_NAMES, op_s, op_s))
-        return :($out)
-    catch
-    end
-    return quote
-        op_s = typeof(op) <: Broadcast.BroadcastFunction ? string(op.f) * '.' : string(op)
-        if length(op_s) == 2 && op_s[1] in ('+', '-', '*', '/', '^') && op_s[2] == '.'
-            op_s = '.' * op_s[1]
+        op_s = if op isa Broadcast.BroadcastFunction
+            base_op_s = string(op.f)
+            if length(base_op_s) == 1 && first(base_op_s) in ('+', '-', '*', '/', '^')
+                # Like `.+`
+                string('.', base_op_s)
+            else
+                # Like `cos.`
+                string(base_op_s, '.')
+            end
+        else
+            string(op)
         end
-        out = collect(get(OP_NAMES, op_s, op_s))
-        return out
+        cache[h] = op_s
+        return op_s
+    finally
+        unlock(OP_NAME_CACHE.lock)
     end
 end
 

src/Utils.jl

@@ -2,6 +2,7 @@
 module UtilsModule
 
 using MacroTools: postwalk, @capture, splitdef, combinedef
+using DispatchDoctor: @unstable
 
 # Returns two arrays
 macro return_on_false2(flag, retval, retval2)
@@ -124,7 +125,7 @@ function deprecate_varmap(variable_names, varMap, func_name)
     return variable_names
 end
 
-counttuple(::Type{<:NTuple{N,Any}}) where {N} = N
+@unstable counttuple(::Type{<:NTuple{N,Any}}) where {N} = N
 
 """
     Undefined

test/test_deprecations.jl

@@ -45,14 +45,6 @@ if VERSION >= v"1.9"
     )
 end
 
-# Old usage of evaluation options
-if VERSION >= v"1.9-"
-    ex = Expression(Node{Float64}(; feature=1))
-    @test_logs (:warn, r"The `turbo` and `bumper` keyword arguments are deprecated.*") (ex(
-        randn(Float64, 1, 10), OperatorEnum(); turbo=true
-    ))
-end
-
 # Test deprecated modules
 logs = @capture_err begin
     @eval using DynamicExpressions.EquationModule

test/test_params.jl

@@ -1,4 +1,5 @@
 using DynamicExpressions
+import DynamicExpressions as DE
 
 maximum_residual = 1e-2
 
@@ -26,6 +27,13 @@ maximum_residual = 1e-2
     greater(x, y) = (x > y)
 
     custom_cos(x) = cos(x)^2
+
+    DE.get_op_name(::typeof(safe_log)) = "log"
+    DE.get_op_name(::typeof(safe_log2)) = "log2"
+    DE.get_op_name(::typeof(safe_log10)) = "log10"
+    DE.get_op_name(::typeof(safe_log1p)) = "log1p"
+    DE.get_op_name(::typeof(safe_acosh)) = "acosh"
+    DE.get_op_name(::typeof(safe_sqrt)) = "sqrt"
 end
 
 HEADER_GUARD_TEST_PARAMS = true

test/test_print.jl

@@ -1,5 +1,6 @@
 using Test
 using DynamicExpressions
+import DynamicExpressions as DE
 import Compat: Returns
 
 include("test_params.jl")
@@ -32,8 +33,10 @@ for unaop in [safe_log, safe_log2, safe_log10, safe_log1p, safe_sqrt, safe_acosh
     @test string_tree(minitree, opts) == replace(string(unaop), "safe_" => "") * "(x1)"
 end
 
-!(@isdefined safe_pow) &&
-    @eval safe_pow(x::T, y::T) where {T<:Number} = (x < 0 && y != round(y)) ? T(NaN) : x^y
+@isdefined(safe_pow) || @eval begin
+    safe_pow(x::T, y::T) where {T<:Number} = (x < 0 && y != round(y)) ? T(NaN) : x^y
+    DE.get_op_name(::typeof(safe_pow)) = "^"
+end
 for binop in [safe_pow, ^]
     opts = OperatorEnum(;
         default_params..., binary_operators=(+, *, /, -, binop), unary_operators=(cos,)