Merge pull request #37 from c42f/cjf/fix-distributed-deserialize

Ensure RGF body is added to the cache when deserializing

Author: ChrisRackauckas

src/RuntimeGeneratedFunctions.jl

@@ -62,6 +62,11 @@ struct RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id} <: Functio
         cached_body = _cache_body(cache_tag, id, body)
         new{Tuple(args), cache_tag, context_tag, id}(cached_body)
     end
+
+    # For internal use in deserialize() - doesen't check whether the body is in the cache!
+    function RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}(body::Expr) where {argnames,cache_tag,context_tag,id}
+        new{argnames, cache_tag, context_tag, id}(body)
+    end
 end
 
 function _check_rgf_initialized(mods...)
@@ -273,6 +278,15 @@ function closures_to_opaque(ex::Expr, return_type=nothing)
     return Expr(head, Any[closures_to_opaque(x, return_type) for x in args]...)
 end
 
+# We write an explicit deserialize() here to trigger caching of the body on a
+# remote node when using Serialialization.jl (in Distributed.jl and elsewhere)
+function Serialization.deserialize(s::AbstractSerializer,
+        ::Type{RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}}) where {argnames,cache_tag,context_tag,id}
+    body = deserialize(s)
+    cached_body = _cache_body(cache_tag, id, body)
+    RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}(cached_body)
+end
+
 @specialize
 
 end

test/runtests.jl

@@ -1,4 +1,5 @@
 using RuntimeGeneratedFunctions, BenchmarkTools
+using Serialization
 using Test
 
 RuntimeGeneratedFunctions.init(@__MODULE__)
@@ -158,3 +159,9 @@ if VERSION >= v"1.7.0-DEV.351"
     ex = :(x -> [2i for i in 1:x])
     @test @RuntimeGeneratedFunction(ex)(3) == [2, 4, 6]
 end
+
+# Serialization
+
+buf = IOBuffer(read(`$(Base.julia_cmd()) "serialize_rgf.jl"`))
+deserialized_f = deserialize(buf)
+@test deserialized_f(11) == "Hi from a separate process. x=11"

test/serialize_rgf.jl

@@ -0,0 +1,10 @@
+# Must be run in a separate process from the rest of the tests!
+
+using RuntimeGeneratedFunctions
+using Serialization
+
+RuntimeGeneratedFunctions.init(@__MODULE__)
+
+f = @RuntimeGeneratedFunction(:(x->"Hi from a separate process. x=$x"))
+
+serialize(stdout, f)