Merge pull request #6 from c42f/cjf/support-precompilation

Support for precompilation

Author: ChrisRackauckas

README.md

@@ -4,8 +4,7 @@
 
 `RuntimeGeneratedFunctions` are functions generated at runtime without world-age
 issues and with the full performance of a standard Julia anonymous function. This
-builds functions in a way that avoids `eval`, but cannot store the precompiled
-functions between Julia sessions.
+builds functions in a way that avoids `eval`.
 
 Note that `RuntimeGeneratedFunction` does not handle closures. Please use the
 [GeneralizedGenerated.jl](https://github.com/JuliaStaging/GeneralizedGenerated.jl)
@@ -23,7 +22,7 @@ function no_worldage()
         @inbounds _du[2] = _u[2]
         nothing
     end)
-    f1 = RuntimeGeneratedFunction(ex)
+    f1 = @RuntimeGeneratedFunction(ex)
     du = rand(2)
     u = rand(2)
     p = nothing

src/RuntimeGeneratedFunctions.jl

@@ -2,27 +2,106 @@ module RuntimeGeneratedFunctions
 
 using ExprTools, Serialization, SHA
 
-const function_cache = Dict{Tuple,Expr}()
-struct RuntimeGeneratedFunction{uuid,argnames}
-    function RuntimeGeneratedFunction(ex)
+export @RuntimeGeneratedFunction
+
+
+"""
+    RuntimeGeneratedFunction
+
+This type should be constructed via the macro @RuntimeGeneratedFunction.
+"""
+struct RuntimeGeneratedFunction{moduletag,id,argnames}
+    function RuntimeGeneratedFunction(moduletag, ex)
         def = splitdef(ex)
         args, body = normalize_args(def[:args]), def[:body]
-        uuid = expr2bytes(body)
-        function_cache[uuid] = body
-        new{uuid,Tuple(args)}()
+        id = expr2bytes(body)
+        _cache_body(moduletag, id, body)
+        new{moduletag,id,Tuple(args)}()
+    end
+end
+
+"""
+    @RuntimeGeneratedFunction(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:
+
+* The result can be called immediately (immune to world age errors)
+* The result is not a named generic function, and doesn't participate in
+  generic function dispatch; it's more like a callable method.
+
+# Examples
+```
+function foo()
+    expression = :((x,y)->x+y+1) # May be generated dynamically
+    f = @RuntimeGeneratedFunction(expression)
+    f(1,2) # May be called immediately
+end
+```
+"""
+macro RuntimeGeneratedFunction(ex)
+    _ensure_cache_exists!(__module__)
+    quote
+        RuntimeGeneratedFunction(
+            $(esc(_tagname)),
+            $(esc(ex))
+        )
     end
 end
+
+function Base.show(io::IO, f::RuntimeGeneratedFunction{moduletag, id, argnames}) where {moduletag,id,argnames}
+    body = _lookup_body(moduletag, id)
+    mod = parentmodule(moduletag)
+    func_expr = Expr(:->, Expr(:tuple, argnames...), body)
+    print(io, "RuntimeGeneratedFunction(#=in $mod=#, ", repr(func_expr), ")")
+end
+
 (f::RuntimeGeneratedFunction)(args::Vararg{Any,N}) where N = generated_callfunc(f, args...)
 
-@inline @generated function generated_callfunc(f::RuntimeGeneratedFunction{uuid,argnames},__args...) where {uuid,argnames}
+@inline @generated function generated_callfunc(f::RuntimeGeneratedFunction{moduletag, id, argnames}, __args...) where {moduletag,id,argnames}
     setup = (:($(argnames[i]) = @inbounds __args[$i]) for i in 1:length(argnames))
     quote
         $(setup...)
-        $(function_cache[uuid])
+        $(_lookup_body(moduletag, id))
     end
 end
 
-export RuntimeGeneratedFunction
+### Function body caching and lookup
+#
+# Caching the body of a RuntimeGeneratedFunction is a little complicated
+# because we want the `id=>body` mapping to survive precompilation. This means
+# we need to store the cache of mappings which are created by a module in that
+# module itself.
+#
+# For that, we need a way to lookup the correct module from an instance of
+# RuntimeGeneratedFunction.  Modules can't be type parameters, but we can use
+# any type which belongs to the module as a proxy "tag" for the module.
+#
+# (We could even abuse `typeof(__module__.eval)` for the tag, though this is a
+# little non-robust to weird special cases like Main.eval being
+# Base.MainInclude.eval.)
+
+_cachename = Symbol("#_RuntimeGeneratedFunctions_cache")
+_tagname = Symbol("#_RuntimeGeneratedFunctions_ModTag")
+
+function _cache_body(moduletag, id, body)
+    getfield(parentmodule(moduletag), _cachename)[id] = body
+end
+
+function _lookup_body(moduletag, id)
+    getfield(parentmodule(moduletag), _cachename)[id]
+end
+
+function _ensure_cache_exists!(mod)
+    if !isdefined(mod, _cachename)
+        mod.eval(quote
+            const $_cachename = Dict{Tuple,Expr}()
+            struct $_tagname
+            end
+        end)
+    end
+end
 
 ###
 ### Utilities

test/precomp/RGFPrecompTest.jl

@@ -0,0 +1,5 @@
+module RGFPrecompTest
+    using RuntimeGeneratedFunctions
+
+    f = @RuntimeGeneratedFunction(:((x,y)->x+y))
+end

test/runtests.jl

@@ -25,9 +25,9 @@ ex3 = :(function (_du::T,_u::Vector{E},_p::P,_t::Any) where {T<:Vector,E,P}
     nothing
 end)
 
-f1 = RuntimeGeneratedFunction(ex1)
-f2 = RuntimeGeneratedFunction(ex2)
-f3 = RuntimeGeneratedFunction(ex3)
+f1 = @RuntimeGeneratedFunction(ex1)
+f2 = @RuntimeGeneratedFunction(ex2)
+f3 = @RuntimeGeneratedFunction(ex3)
 
 du = rand(2)
 u = rand(2)
@@ -58,11 +58,25 @@ function no_worldage()
         @inbounds _du[2] = _u[2]
         nothing
     end)
-    f1 = RuntimeGeneratedFunction(ex)
+    f1 = @RuntimeGeneratedFunction(ex)
     du = rand(2)
     u = rand(2)
     p = nothing
     t = nothing
     f1(du,u,p,t)
 end
 @test no_worldage() === nothing
+
+# Test show()
+@test sprint(show, @RuntimeGeneratedFunction(Base.remove_linenums!(:((x,y)->x+y+1)))) ==
+     """
+     RuntimeGeneratedFunction(#=in $(@__MODULE__)=#, :((x, y)->begin
+               x + y + 1
+           end))"""
+
+# Test with precompilation
+push!(LOAD_PATH, joinpath(@__DIR__, "precomp"))
+using RGFPrecompTest
+
+@test RGFPrecompTest.f(1,2) == 3
+