Restore dropped RGF serialization (#138)

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>

Author: ChrisRackauckas-Claude

src/RuntimeGeneratedFunctions.jl

@@ -2,7 +2,7 @@ module RuntimeGeneratedFunctions
 
 using ExprTools: ExprTools, combinedef, splitdef
 using SHA: SHA, SHA1_CTX, update!
-using Serialization: Serialization, AbstractSerializer, deserialize
+using Serialization: Serialization, AbstractSerializer, deserialize, serialize
 
 export RuntimeGeneratedFunction, @RuntimeGeneratedFunction, drop_expr
 
@@ -49,10 +49,9 @@ it represents one callable method and does not participate in method dispatch.
 
 # Serialization
 
-`Serialization.serialize` preserves a runtime-generated function while its
-`body` is retained. Serialize a function before calling [`drop_expr`](@ref): a
-dropped function intentionally does not retain the expression needed to restore
-its cache in another process.
+`Serialization.serialize` preserves a runtime-generated function, including
+after [`drop_expr`](@ref). A dropped function's expression is recovered from its
+cache while serializing and remains dropped after deserialization.
 
 # Examples
 
@@ -112,9 +111,6 @@ function callable.
 The expression can still be retrieved later using [`get_expression`](@ref) as long
 as at least one `RuntimeGeneratedFunction` with the same body exists.
 
-Serialize a generated function before calling `drop_expr`. A dropped function has
-discarded the expression needed to restore its cache in another process.
-
 # Examples
 ```julia
 ex = :((x) -> x^2)
@@ -511,6 +507,51 @@ function get_expression(
     return func_expr = Expr(:->, Expr(:tuple, argnames...), _lookup_body(cache_tag, id))
 end
 
+struct _SerializedRuntimeGeneratedFunction{argnames, cache_tag, context_tag, id, B}
+    body::B
+end
+
+function Serialization.serialize(
+        s::AbstractSerializer,
+        ::RuntimeGeneratedFunction{
+            argnames, cache_tag,
+            context_tag, id, Nothing,
+        }
+    ) where {
+        argnames,
+        cache_tag,
+        context_tag,
+        id,
+    }
+    body = _lookup_body(cache_tag, id)
+    proxy = _SerializedRuntimeGeneratedFunction{
+        argnames, cache_tag, context_tag, id, typeof(body),
+    }(body)
+    return serialize(s, proxy)
+end
+
+function Serialization.deserialize(
+        s::AbstractSerializer,
+        ::Type{
+            _SerializedRuntimeGeneratedFunction{
+                argnames, cache_tag,
+                context_tag, id, B,
+            },
+        }
+    ) where {
+        argnames,
+        cache_tag,
+        context_tag,
+        id,
+        B,
+    }
+    body = deserialize(s)
+    cached_body = _cache_body(cache_tag, id, body)
+    return drop_expr(
+        RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id}(cached_body)
+    )
+end
+
 function Serialization.deserialize(
         s::AbstractSerializer,
         ::Type{

test/core_tests.jl

@@ -212,14 +212,11 @@ proj = dirname(Base.active_project())
 serialize_script = joinpath(@__DIR__, "shared", "serialize_rgf.jl")
 julia = joinpath(Sys.BINDIR, "julia")
 buf = IOBuffer(read(`$julia --startup-file=no --project=$proj $serialize_script`))
-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
-
-serialization_buffer = IOBuffer()
-serialize(serialization_buffer, drop_expr(@RuntimeGeneratedFunction(:(x -> x + 1))))
-seekstart(serialization_buffer)
-@test_throws ArgumentError deserialize(serialization_buffer)
+@test deserialized_g(12) == "Serialization with dropped body. y=12"
+@test deserialized_g.body isa Nothing
 
 # deepcopy
 ff = @RuntimeGeneratedFunction(:(x -> [x, x + 1]))

test/shared/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))