Revert "Remove extra precompilation"

This reverts commit 973e3e6.

Project.toml

@@ -16,6 +16,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -46,6 +47,7 @@ MacroTools = "0.4, 0.5"
 Optim = "0.19, 1.1 - 1.7.6"
 PackageExtensionCompat = "1"
 Pkg = "1"
+PrecompileTools = "1"
 ProgressBars = "1.4"
 Reexport = "1"
 SpecialFunctions = "0.10.1, 1, 2"

src/SymbolicRegression.jl

@@ -1136,4 +1136,11 @@ macro ignore(args...) end
 # Hack to get static analysis to work from within tests:
 @ignore include("../test/runtests.jl")
 
+include("precompile.jl")
+redirect_stdout(devnull) do
+    redirect_stderr(devnull) do
+        do_precompilation(Val(:precompile))
+    end
+end
+
 end #module SR

src/precompile.jl

@@ -0,0 +1,87 @@
+import PrecompileTools: @compile_workload, @setup_workload
+
+macro maybe_setup_workload(mode, ex)
+    precompile_ex = Expr(
+        :macrocall, Symbol("@setup_workload"), LineNumberNode(@__LINE__), ex
+    )
+    return quote
+        if $(esc(mode)) == :compile
+            $(esc(ex))
+        elseif $(esc(mode)) == :precompile
+            $(esc(precompile_ex))
+        else
+            error("Invalid value for mode: " * show($(esc(mode))))
+        end
+    end
+end
+
+macro maybe_compile_workload(mode, ex)
+    precompile_ex = Expr(
+        :macrocall, Symbol("@compile_workload"), LineNumberNode(@__LINE__), ex
+    )
+    return quote
+        if $(esc(mode)) == :compile
+            $(esc(ex))
+        elseif $(esc(mode)) == :precompile
+            $(esc(precompile_ex))
+        else
+            error("Invalid value for mode: " * show($(esc(mode))))
+        end
+    end
+end
+
+"""`mode=:precompile` will use `@precompile_*` directives; `mode=:compile` runs."""
+function do_precompilation(::Val{mode}) where {mode}
+    @maybe_setup_workload mode begin
+        for T in [Float32, Float64], nout in 1:2
+            start = nout == 1
+            N = 30
+            X = randn(T, 3, N)
+            y = start ? randn(T, N) : randn(T, nout, N)
+            @maybe_compile_workload mode begin
+                options = SymbolicRegression.Options(;
+                    binary_operators=[+, *, /, -, ^],
+                    unary_operators=[sin, cos, exp, log, sqrt, abs],
+                    populations=3,
+                    population_size=start ? 50 : 12,
+                    ncycles_per_iteration=start ? 30 : 1,
+                    mutation_weights=MutationWeights(;
+                        mutate_constant=1.0,
+                        mutate_operator=1.0,
+                        add_node=1.0,
+                        insert_node=1.0,
+                        delete_node=1.0,
+                        simplify=1.0,
+                        randomize=1.0,
+                        do_nothing=1.0,
+                        optimize=1.0,
+                    ),
+                    fraction_replaced=0.2,
+                    fraction_replaced_hof=0.2,
+                    define_helper_functions=false,
+                    optimizer_probability=0.05,
+                    save_to_file=false,
+                )
+                state = equation_search(
+                    X,
+                    y;
+                    niterations=start ? 3 : 1,
+                    options=options,
+                    parallelism=:multithreading,
+                    return_state=true,
+                )
+                hof = equation_search(
+                    X,
+                    y;
+                    niterations=0,
+                    options=options,
+                    parallelism=:multithreading,
+                    saved_state=state,
+                    return_state=false,
+                )
+                nout == 1 && calculate_pareto_frontier(hof)
+            end
+        end
+    end
+    return nothing
+end

test/full.jl

@@ -40,6 +40,10 @@ end
     include("test_custom_operators_multiprocessing.jl")
 end
 
+@testset "Test whether the precompilation script works." begin
+    include("test_precompilation.jl")
+end
+
 @testset "Test whether custom objectives work." begin
     include("test_custom_objectives.jl")
 end

test/test_precompilation.jl

@@ -0,0 +1,3 @@
+using SymbolicRegression
+
+SymbolicRegression.do_precompilation(Val(:compile))