From 235adbeef33203e0cd8cd431f38f652d8c2e4a60 Mon Sep 17 00:00:00 2001
From: Simeon Schaub <simeondavidschaub99@gmail.com>
Date: Tue, 6 Apr 2021 11:21:36 +0200
Subject: [PATCH] support for closures via opaque closures

Should close #28. Kwargs or `where` parameters might not work, but that's an issue in Julia base.
---
 src/RuntimeGeneratedFunctions.jl | 67 ++++++++++++++++++++++++++++++--
 test/runtests.jl                 | 24 ++++++++++++
 2 files changed, 87 insertions(+), 4 deletions(-)

diff --git a/src/RuntimeGeneratedFunctions.jl b/src/RuntimeGeneratedFunctions.jl
index fae241b..11da4ba 100644
--- a/src/RuntimeGeneratedFunctions.jl
+++ b/src/RuntimeGeneratedFunctions.jl
@@ -11,9 +11,13 @@ export @RuntimeGeneratedFunction
 This type should be constructed via the macro @RuntimeGeneratedFunction.
 """
 struct RuntimeGeneratedFunction{moduletag,id,argnames}
-    function RuntimeGeneratedFunction(moduletag, ex)
+    function RuntimeGeneratedFunction(moduletag, ex; opaque_closures=true)
         def = splitdef(ex)
         args, body = normalize_args(def[:args]), def[:body]
+        if opaque_closures && isdefined(Base, :Experimental) &&
+            isdefined(Base.Experimental, Symbol("@opaque"))
+            body = closures_to_opaque(body)
+        end
         id = expr2bytes(body)
         _cache_body(moduletag, id, body)
         new{moduletag,id,Tuple(args)}()
@@ -21,7 +25,7 @@ struct RuntimeGeneratedFunction{moduletag,id,argnames}
 end
 
 """
-    @RuntimeGeneratedFunction(function_expression)
+    @RuntimeGeneratedFunction(function_expression, opaque_closures=true)
 
 Construct a function from `function_expression` which can be called immediately
 without world age problems. Somewhat like using `eval(function_expression)` and
@@ -31,6 +35,13 @@ then calling the resulting function. The differences are:
 * The result is not a named generic function, and doesn't participate in
   generic function dispatch; it's more like a callable method.
 
+If `opaque_closures` is `true`, all closures in `function_expression` are
+converted to
+[opaque closures](https://github.com/JuliaLang/julia/pull/37849#issue-496641229).
+This allows for the use of closures and generators inside the generated function,
+but may not work in all cases due to slightly different semantics. This feature
+requires Julia 1.7.
+
 # Examples
 ```
 function foo()
@@ -40,12 +51,13 @@ function foo()
 end
 ```
 """
-macro RuntimeGeneratedFunction(ex)
+macro RuntimeGeneratedFunction(ex, opaque_closures=true)
     _ensure_cache_exists!(__module__)
     quote
         RuntimeGeneratedFunction(
             $(esc(_tagname)),
-            $(esc(ex))
+            $(esc(ex));
+            opaque_closures = $opaque_closures,
         )
     end
 end
@@ -119,4 +131,51 @@ function expr2bytes(ex)
     return Tuple(sha512(take!(io)))
 end
 
+@nospecialize
+
+closures_to_opaque(x, _=nothing) = x
+_tconvert(T, x) = Expr(:(::), Expr(:call, GlobalRef(Base, :convert), T, x), T)
+function closures_to_opaque(ex::Expr, return_type=nothing)
+    head, args = ex.head, ex.args
+    fdef = splitdef(ex; throw=false)
+    if fdef !== nothing
+        body = get(fdef, :body, nothing)
+        if haskey(fdef, :rtype)
+            body = _tconvert(fdef[:rtype], closures_to_opaque(body, fdef[:rtype]))
+            delete!(fdef, :rtype)
+        else
+            body = closures_to_opaque(body)
+        end
+        fdef[:head] = :(->)
+        fdef[:body] = body
+        name = get(fdef, :name, nothing)
+        name !== nothing && delete!(fdef, :name)
+        _ex = Expr(:opaque_closure, combinedef(fdef))
+        # TODO: emit named opaque closure for better stacktraces
+        # (ref https://github.com/JuliaLang/julia/pull/40242)
+        if name !== nothing
+            name isa Symbol ||
+                error("Unsupported function definition `$ex` in RuntimeGeneratedFunction.")
+            _ex = Expr(:(=), name, _ex)
+        end
+        return _ex
+    elseif head === :generator
+        f_args = Expr(:tuple, Any[x.args[1] for x in args[2:end]]...)
+        iters = Any[x.args[2] for x in args[2:end]]
+        return Expr(
+            :call,
+            GlobalRef(Base, :Generator),
+            closures_to_opaque(Expr(:(->), f_args, args[1])),
+            iters...,
+        )
+    elseif head === :opaque_closure
+        return closures_to_opaque(args[1])
+    elseif head === :return && return_type !== nothing
+        return Expr(:return, _tconvert(return_type, closures_to_opaque(args[1], return_type)))
+    end
+    return Expr(head, Any[closures_to_opaque(x, return_type) for x in args]...)
+end
+
+@specialize
+
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index ae06761..d02ee90 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -80,3 +80,27 @@ using RGFPrecompTest
 
 @test RGFPrecompTest.f(1,2) == 3
 
+# closures
+if VERSION >= v"1.7.0-DEV.351"
+    ex = :(x -> (y -> x + y))
+    @test @RuntimeGeneratedFunction(ex)(2)(3) === 5
+
+    ex = :(x -> (f(y::Int)::Float64 = x + y; f))
+    @test @RuntimeGeneratedFunction(ex)(2)(3) === 5.0
+
+    ex = :(x -> function (y::Int)
+                    return x + y
+                end)
+    @test @RuntimeGeneratedFunction(ex)(2)(3) === 5
+
+    ex = :(x -> function f(y::Int)::UInt8
+                    return x + y
+                end)
+    @test @RuntimeGeneratedFunction(ex)(2)(3) === 0x05
+
+    ex = :(x -> sum(i^2 for i in 1:x))
+    @test @RuntimeGeneratedFunction(ex)(3) === 14
+
+    ex = :(x -> [2i for i in 1:x])
+    @test @RuntimeGeneratedFunction(ex)(3) == [2, 4, 6]
+end