Fix serialization of RGFs with dropped body expressions

Author: c42f

src/RuntimeGeneratedFunctions.jl

@@ -319,20 +319,31 @@ 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)
+# We write an explicit serialize() and deserialize() here to manage caching of
+# the body on a remote node when using Serialialization.jl (in Distributed.jl
+# and elsewhere)
+function Serialization.serialize(s::AbstractSerializer,
+        rgf::RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id, B}) where {
+            argnames, cache_tag, context_tag, id, B}
+    body = _lookup_body(cache_tag, id)
+    Serialization.serialize_type(s, RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id, B})
+    serialize(s, body)
+end
+
 function Serialization.deserialize(s::AbstractSerializer,
                                    ::Type{
                                           <:RuntimeGeneratedFunction{argnames, cache_tag,
-                                                                     context_tag, id}}) where {
+                                                                     context_tag, id, B}}) where {
                                                                                                argnames,
                                                                                                cache_tag,
                                                                                                context_tag,
-                                                                                               id
+                                                                                               id,
+                                                                                               B
                                                                                                }
     body = deserialize(s)
     cached_body = _cache_body(cache_tag, id, body)
-    RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}(cached_body)
+    f = RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}(cached_body)
+    B === Nothing ? drop_expr(f) : f
 end
 
 @specialize

test/runtests.jl

@@ -177,5 +177,8 @@ end
 
 proj = dirname(Base.active_project())
 buf = IOBuffer(read(`$(Base.julia_cmd()) --startup-file=no --project=$proj "serialize_rgf.jl"`))
-deserialized_f = deserialize(buf)
+deserialized_f, deserialized_g = deserialize(buf)
 @test deserialized_f(11) == "Hi from a separate process. x=11"
+@test deserialized_f.body isa Expr
+@test deserialized_g(12) == "Serialization with dropped body. y=12"
+@test deserialized_g.body isa Nothing

test/serialize_rgf.jl

@@ -6,5 +6,6 @@ using Serialization
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
 f = @RuntimeGeneratedFunction(:(x -> "Hi from a separate process. x=$x"))
+g = drop_expr(@RuntimeGeneratedFunction(:(y -> "Serialization with dropped body. y=$y")))
 
-serialize(stdout, f)
+serialize(stdout, (f,g))