Merge pull request #26 from c42f/cjf/cache-module-in-constructor

Allow cache module to be given in RGF constructor

Project.toml

@@ -1,7 +1,7 @@
 name = "RuntimeGeneratedFunctions"
 uuid = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 authors = ["Chris Rackauckas <[email protected]> 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.