Work around to specify cache module when building RuntimeGeneratedFunctions

src/build_function.jl

@@ -101,7 +101,10 @@ function _build_and_inject_function(mod::Module, ex)
     elseif ex.head == :(->)
         return _build_and_inject_function(mod, Expr(:function, ex.args...))
     end
-    @RuntimeGeneratedFunction(mod, ex)
+    # XXX: Workaround to specify the module as both the cache module AND context module.
+    # Currently, the @RuntimeGeneratedFunction macro only sets the context module.
+    module_tag = getproperty(mod, RuntimeGeneratedFunctions._tagname)
+    RuntimeGeneratedFunctions.RuntimeGeneratedFunction(module_tag, module_tag, ex)
 end
 
 # Detect heterogeneous element types of "arrays of matrices/sparce matrices"

test/precompile_test.jl

@@ -8,19 +8,15 @@ using ODEPrecompileTest
 u  = collect(1:3)
 p  = collect(4:6)
 
-# This case does not work, because "f_bad" gets defined in ModelingToolkit
-# instead of in the compiled module!
+# These cases do not work, because they get defined in the ModelingToolkit's RGF cache.
 @test parentmodule(typeof(ODEPrecompileTest.f_bad.f.f_iip).parameters[2]) == ModelingToolkit
 @test parentmodule(typeof(ODEPrecompileTest.f_bad.f.f_oop).parameters[2]) == ModelingToolkit
 @test parentmodule(typeof(ODEPrecompileTest.f_noeval_bad.f.f_iip).parameters[2]) == ModelingToolkit
 @test parentmodule(typeof(ODEPrecompileTest.f_noeval_bad.f.f_oop).parameters[2]) == ModelingToolkit
 @test_throws KeyError ODEPrecompileTest.f_bad(u, p, 0.1)
 @test_throws KeyError ODEPrecompileTest.f_noeval_bad(u, p, 0.1)
 
-# This case works, because "f_good" gets defined in the precompiled module.
-@test parentmodule(typeof(ODEPrecompileTest.f_good.f.f_iip).parameters[2]) == ODEPrecompileTest
-@test parentmodule(typeof(ODEPrecompileTest.f_good.f.f_oop).parameters[2]) == ODEPrecompileTest
+# This case works, because it gets defined with the appropriate cache and context tags.
 @test parentmodule(typeof(ODEPrecompileTest.f_noeval_good.f.f_iip).parameters[2]) == ODEPrecompileTest
 @test parentmodule(typeof(ODEPrecompileTest.f_noeval_good.f.f_oop).parameters[2]) == ODEPrecompileTest
-@test ODEPrecompileTest.f_good(u, p, 0.1) == [4, 0, -16]
 @test ODEPrecompileTest.f_noeval_good(u, p, 0.1) == [4, 0, -16]

test/precompile_test/ODEPrecompileTest.jl

@@ -5,7 +5,7 @@ module ODEPrecompileTest
         # Define some variables
         @parameters t σ ρ β
         @variables x(t) y(t) z(t)
-        @derivatives D'~t
+        D = Differential(t)
 
         # Define a differential equation
         eqs = [D(x) ~ σ*(y-x),
@@ -16,17 +16,14 @@ module ODEPrecompileTest
         return ODEFunction(de, [x,y,z], [σ,ρ,β]; kwargs...)
     end
 
-    # Build an ODEFunction as part of the module's precompilation. This case
-    # will not work, because the generated RGFs will be put into
-    # ModelingToolkit's RGF cache.
+    # Build an ODEFunction as part of the module's precompilation. These cases
+    # will not work, because the generated RGFs are put into the ModelingToolkit cache.
     const f_bad = system()
+    const f_noeval_bad = system(; eval_expression=false)
 
-    # This case will work, because it will be put into our own module's cache.
+    # Setting eval_expression=false and eval_module=[this module] will ensure
+    # the RGFs are put into our own cache, initialised below.
     using RuntimeGeneratedFunctions
     RuntimeGeneratedFunctions.init(@__MODULE__)
-    const f_good = system(; eval_module=@__MODULE__)
-
-    # Also test that eval_expression=false works
-    const f_noeval_bad = system(; eval_expression=false)
     const f_noeval_good = system(; eval_expression=false, eval_module=@__MODULE__)
 end

test/runtests.jl

@@ -29,9 +29,9 @@ using SafeTestsets, Test
 @safetestset "Depdendency Graph Test" begin include("dep_graphs.jl") end
 @safetestset "Function Registration Test" begin include("function_registration.jl") end
 @safetestset "Array of Array Test" begin include("build_function_arrayofarray.jl") end
+@safetestset "Precompiled Modules Test" begin include("precompile_test.jl") end
 @testset "Distributed Test" begin include("distributed.jl") end
 @safetestset "Variable Utils Test" begin include("variable_utils.jl") end
 println("Last test requires gcc available in the path!")
 @safetestset "C Compilation Test" begin include("ccompile.jl") end
 @safetestset "Latexify recipes Test" begin include("latexify.jl") end
-@safetestset "Precompiled Modules Test" begin include("precompile_test.jl") end