Allow cache module to be given in RGF constructor

This exposes the full constructor (without the need for specifying
module tags), as we seem to need this for utility functions which want
to cache the RGF in a user-specified module.

Author: c42f

Project.toml

@@ -1,7 +1,7 @@
 name = "RuntimeGeneratedFunctions"
 uuid = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"

src/RuntimeGeneratedFunctions.jl

@@ -2,29 +2,15 @@ module RuntimeGeneratedFunctions
 
 using ExprTools, Serialization, SHA
 
-export @RuntimeGeneratedFunction
+export RuntimeGeneratedFunction, @RuntimeGeneratedFunction
 
 
-"""
-    RuntimeGeneratedFunction
-
-This type should be constructed via the macro @RuntimeGeneratedFunction.
-"""
-struct RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id} <: Function
-    body::Expr
-    function RuntimeGeneratedFunction(cache_tag, context_tag, ex)
-        def = splitdef(ex)
-        args, body = normalize_args(def[:args]), def[:body]
-        id = expr_to_id(body)
-        cached_body = _cache_body(cache_tag, id, body)
-        new{Tuple(args), cache_tag, context_tag, id}(cached_body)
-    end
-end
-
 """
     @RuntimeGeneratedFunction(function_expression)
     @RuntimeGeneratedFunction(context_module, function_expression)
 
+    RuntimeGeneratedFunction(cache_module, context_module, function_expression)
+
 Construct a function from `function_expression` which can be called immediately
 without world age problems. Somewhat like using `eval(function_expression)` and
 then calling the resulting function. The differences are:
@@ -40,6 +26,11 @@ If provided, `context_module` is module in which symbols within
 `function_expression` will be looked up. By default this is module in which
 `@RuntimeGeneratedFunction` is expanded.
 
+`cache_module` is the module where the expression `code` will be cached. If
+`RuntimeGeneratedFunction` is used during precompilation, this must be a module
+which is currently being precompiled. Normally this would be set to
+`@__MODULE__` using one of the macro constructors.
+
 # Examples
 ```
 RuntimeGeneratedFunctions.init(@__MODULE__) # Required at module top-level
@@ -51,28 +42,46 @@ function foo()
 end
 ```
 """
-macro RuntimeGeneratedFunction(code)
-    _RGF_constructor_code(:(@__MODULE__), esc(code))
-end
-macro RuntimeGeneratedFunction(context_module, code)
-    _RGF_constructor_code(esc(context_module), esc(code))
+struct RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id} <: Function
+    body::Expr
+    function RuntimeGeneratedFunction(cache_tag, context_tag, ex)
+        def = splitdef(ex)
+        args, body = normalize_args(def[:args]), def[:body]
+        id = expr_to_id(body)
+        cached_body = _cache_body(cache_tag, id, body)
+        new{Tuple(args), cache_tag, context_tag, id}(cached_body)
+    end
 end
 
-function _RGF_constructor_code(context_module, code)
-    quote
-        code = $code
-        cache_module = @__MODULE__
-        context_module = $context_module
-        if #==# !isdefined(cache_module,   $(QuoteNode(_tagname))) ||
-                !isdefined(context_module, $(QuoteNode(_tagname)))
-            init_mods = unique([context_module, cache_module])
+function _check_rgf_initialized(mods...)
+    for mod in mods
+        if !isdefined(mod, _tagname)
             error("""You must use `RuntimeGeneratedFunctions.init(@__MODULE__)` at module
-                     top level before using runtime generated functions in $init_mods""")
+                  top level before using runtime generated functions in $mod""")
         end
-        RuntimeGeneratedFunction(cache_module.$_tagname, context_module.$_tagname, $code)
     end
 end
 
+function RuntimeGeneratedFunction(cache_module::Module, context_module::Module, code)
+    _check_rgf_initialized(cache_module, context_module)
+    RuntimeGeneratedFunction(getfield(cache_module, _tagname),
+                             getfield(context_module, _tagname), code)
+end
+
+macro RuntimeGeneratedFunction(code)
+    quote
+        RuntimeGeneratedFunction(@__MODULE__, @__MODULE__, $(esc(code)))
+    end
+end
+macro RuntimeGeneratedFunction(context_module, code)
+    quote
+        RuntimeGeneratedFunction(@__MODULE__, $(esc(context_module)), $(esc(code)))
+    end
+end
+
+# Duplicate RuntimeGeneratedFunction docs onto @RuntimeGeneratedFunction
+@eval @doc $(@doc RuntimeGeneratedFunction) var"@RuntimeGeneratedFunction"
+
 function Base.show(io::IO, ::MIME"text/plain", f::RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}) where {argnames,cache_tag,context_tag,id}
     cache_mod = parentmodule(cache_tag)
     context_mod = parentmodule(context_tag)

test/precomp/RGFPrecompTest.jl

@@ -1,6 +1,9 @@
 module RGFPrecompTest
     using RuntimeGeneratedFunctions
+    using RGFPrecompTest2
     RuntimeGeneratedFunctions.init(@__MODULE__)
 
     f = @RuntimeGeneratedFunction(:((x,y)->x+y))
+
+    g = RGFPrecompTest2.generate_rgf(@__MODULE__)
 end

test/precomp/RGFPrecompTest2.jl

@@ -0,0 +1,13 @@
+module RGFPrecompTest2
+    using RuntimeGeneratedFunctions
+    RuntimeGeneratedFunctions.init(@__MODULE__)
+
+    y_in_RGFPrecompTest2 = 2
+
+    # Simulates a helper function which generates an RGF, but caches it in a
+    # different module.
+    function generate_rgf(cache_module)
+        context_module = @__MODULE__
+        RuntimeGeneratedFunction(cache_module, @__MODULE__, :((x)->y_in_RGFPrecompTest2+x))
+    end
+end

test/runtests.jl

@@ -83,6 +83,7 @@ push!(LOAD_PATH, joinpath(@__DIR__, "precomp"))
 using RGFPrecompTest
 
 @test RGFPrecompTest.f(1,2) == 3
+@test RGFPrecompTest.g(40) == 42
 
 # Test that RuntimeGeneratedFunction with identical body expressions (but
 # allocated separately) don't clobber each other when one is GC'd.