Closures in RuntimeGeneratedFunctions

[0x0f0f0f]

Just encountered an odd bug testing my package

    Reading Memory: Error During Test at /home/sea/src/julia/Metatheory/test/test_while_interpreter.jl:17
      Test threw exception
      Expression: 2 == rewrite(:((x, $(Mem(:x => 2)))), read_mem; order = :inner, m = #= /home/sea/src/julia/Metatheory/test/test_while_interpreter.jl:17 =# @__MODULE__())
      MethodError: no method matching generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(Symbol("##reducing_expression#257"),), var"#_RGF_ModTag", var"#_RGF_ModTag", (0xa3a3c3bf, 0xb070e435, 0x08893b98, 0x9f9e31fb, 0xf15bd3ff)}, ::Symbol, ::Module)
      The applicable method may be too new: running in world age 29639, while current world is 29642.
      Closest candidates are:
        generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{argnames, cache_tag, var"#_RGF_ModTag", id}, ::Any...) where {argnames, cache_tag, id} at none:0 (method too new to be called from this world context.)
        generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{argnames, cache_tag, Metatheory.var"#_RGF_ModTag", id}, ::Any...) where {argnames, cache_tag, id} at none:0
      Stacktrace:
        [1] (::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(Symbol("##reducing_expression#257"),), var"#_RGF_ModTag", var"#_RGF_ModTag", (0xa3a3c3bf, 0xb070e435, 0x08893b98, 0x9f9e31fb, 0xf15bd3ff)})(::Symbol, ::Module)
          @ RuntimeGeneratedFunctions ~/.julia/packages/RuntimeGeneratedFunctions/tJEmP/src/RuntimeGeneratedFunctions.jl:92
        [2] (::Metatheory.var"#35#41"{Module})(x::Symbol)
          @ Metatheory ~/src/julia/Metatheory/src/rewrite.jl:24
        [3] normalize_nocycle(::Function, ::Symbol; callback::Metatheory.var"#34#40"{Int64})
          @ Metatheory ~/src/julia/Metatheory/src/util.jl:119
        [4] #36
          @ ~/src/julia/Metatheory/src/rewrite.jl:25 [inlined]
        [5] #df_walk!#6
          @ ~/src/julia/Metatheory/src/util.jl:30 [inlined]
        [6] #7
          @ ~/src/julia/Metatheory/src/util.jl:38 [inlined]
        [7] |>(x::Symbol, f::Metatheory.var"#7#8"{Vector{Symbol}, Bool, Metatheory.var"#36#42"{Metatheory.var"#35#41"{Module}, Metatheory.var"#34#40"{Int64}}, Tuple{}})
          @ Base ./operators.jl:859
        [8] _broadcast_getindex_evalf
          @ ./broadcast.jl:648 [inlined]
        [9] _broadcast_getindex
          @ ./broadcast.jl:621 [inlined]
       [10] getindex
          @ ./broadcast.jl:575 [inlined]
       [11] copy
          @ ./broadcast.jl:922 [inlined]
       [12] materialize(bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(|>), Tuple{Vector{Any}, Base.RefValue{Metatheory.var"#7#8"{Vector{Symbol}, Bool, Metatheory.var"#36#42"{Metatheory.var"#35#41"{Module}, Metatheory.var"#34#40"{Int64}}, Tuple{}}}}})
          @ Base.Broadcast ./broadcast.jl:883
       [13] df_walk!(::Function, ::Expr; skip::Vector{Symbol}, skip_call::Bool)
          @ Metatheory ~/src/julia/Metatheory/src/util.jl:38
       [14] #37
          @ ~/src/julia/Metatheory/src/rewrite.jl:30 [inlined]
       [15] (::Metatheory.var"#39#45"{Metatheory.var"#37#43", Metatheory.var"#36#42"{Metatheory.var"#35#41"{Module}, Metatheory.var"#34#40"{Int64}}})(x::Expr)
          @ Metatheory ~/src/julia/Metatheory/src/rewrite.jl:37
       [16] normalize_nocycle(::Function, ::Expr; callback::Metatheory.var"#24#26")
          @ Metatheory ~/src/julia/Metatheory/src/util.jl:119
       [17] normalize_nocycle(::Function, ::Expr)
          @ Metatheory ~/src/julia/Metatheory/src/util.jl:117
       [18] rewrite(ex::Expr, theory::Vector{Rule}; __source__::LineNumberNode, order::Symbol, m::Module, timeout::Int64)
          @ Metatheory ~/src/julia/Metatheory/src/rewrite.jl:37
       [19] macro expansion
          @ ~/src/julia/Metatheory/test/test_while_interpreter.jl:17 [inlined]
       [20] macro expansion
          @ ~/src/julia-compiler/usr/share/julia/stdlib/v1.6/Test/src/Test.jl:1151 [inlined]
       [21] top-level scope
          @ ~/src/julia/Metatheory/test/test_while_interpreter.jl:17

How to reproduce: Comment lines 7 and 8 in Metatheory.jl/test/runtests.jl https://github.com/0x0f0f0f/Metatheory.jl/blob/a273d9ed6b7d88cadaa3d4f66299de3f3649d719/test/runtests.jl#L7-L8

uncomment last argument (cache module) in line 14 in Metatheory.jl/test/test_while_interpreter.jl https://github.com/0x0f0f0f/Metatheory.jl/blob/a273d9ed6b7d88cadaa3d4f66299de3f3649d719/test/test_while_interpreter.jl#L14

Am I just initializing/using RGF in the wrong way? (here and also here) I admit that I have done some weird hacks to achieve a behaviour that is partially close to dynamic scoped variable capturing (not technically closures but close to).

[ChrisRackauckas]

This is complex enough that a simpler reproducer will likely be necessary to narrow it down.

[0x0f0f0f]

module Foo
using RuntimeGeneratedFunctions
const RGF = RuntimeGeneratedFunctions
RGF.init(@__MODULE__)
function genclosure(body, mod::Module)
    (mod != @__MODULE__) && !isdefined(mod, RGF._tagname) && RGF.init(mod)
    RuntimeGeneratedFunction(mod, mod, :(x -> $body))
end
function bar(n::Int; mod = @__MODULE__)
    f = genclosure(:(x * $n), mod)
    g = x -> 2 * f(x)
    g(3)
end
export bar
export genclosure

# also appears to error 
macro baz()
    quote
        n -> bar(n, mod = $__module__)
    end
end
end

julia> using .Foo

# OK
julia> bar(3)
18

julia> bar(3, mod=@__MODULE__)
ERROR: MethodError: no method matching generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x,), var"#_RGF_ModTag", var"#_RGF_ModTag", (0x346def3e, 0x769edc98, 0x33949a6f, 0xc294c197, 0xcea7fb6f)}, ::Int64)
The applicable method may be too new: running in world age 29602, while current world is 29605.
Closest candidates are:
  generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{argnames, cache_tag, var"#_RGF_ModTag", id}, ::Any...) where {argnames, cache_tag, id} at none:0 (method too new to be called from this world context.)
  generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{argnames, cache_tag, Main.Foo.var"#_RGF_ModTag", id}, ::Any...) where {argnames, cache_tag, id} at none:0
Stacktrace:
 [1] (::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x,), var"#_RGF_ModTag", var"#_RGF_ModTag", (0x346def3e, 0x769edc98, 0x33949a6f, 0xc294c197, 0xcea7fb6f)})(args::Int64)
   @ RuntimeGeneratedFunctions ~/.julia/packages/RuntimeGeneratedFunctions/tJEmP/src/RuntimeGeneratedFunctions.jl:92
 [2] (::Main.Foo.var"#3#4"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x,), var"#_RGF_ModTag", var"#_RGF_ModTag", (0x346def3e, 0x769edc98, 0x33949a6f, 0xc294c197, 0xcea7fb6f)}})(x::Int64)
   @ Main.Foo ./REPL[1]:11
 [3] bar(n::Int64; mod::Module)
   @ Main.Foo ./REPL[1]:12
 [4] top-level scope
   @ REPL[4]:1

# Calling it again works
julia> bar(3, mod=@__MODULE__)
18

[ChrisRackauckas]

Yeah I don't think internal function definitions are allowed.

[0x0f0f0f]

Im not sure I understood. Do you think there is a workaround?

[ChrisRackauckas]

I'm not sure. Make it a RuntimeGeneratedFunction?

[0x0f0f0f]

Which function should be runtime generated? genclosure or bar?

[ChrisRackauckas]

All? I just don't think it's going to work out because generated functions need purity. But maybe you can fake it like that.

[0x0f0f0f]

Thanks. It worked. It is hacky but really does the job:

module Foo
using RuntimeGeneratedFunctions
const RGF = RuntimeGeneratedFunctions
RGF.init(@__MODULE__)

function closure_gen(mod::Module)
    RuntimeGeneratedFunction((@__MODULE__), (@__MODULE__),
        :( body -> begin
            ($mod != @__MODULE__) && !isdefined($mod, RGF._tagname) && RGF.init($mod)
            RuntimeGeneratedFunction($mod, $mod, :(x -> $body))
        end
        ))
    end

function barr(n::Int; mod = @__MODULE__)
    f = (closure_gen(mod))(:(x * $n))
    g = x -> 2 * f(x)
    g(3)
end

export barr

end

[ChrisRackauckas]

Great!

[0x0f0f0f]

Here's a generic version:

function closure_generator(mod::Module)
    RuntimeGeneratedFunction((@__MODULE__), (@__MODULE__),
    :( expr -> begin
        ($mod != @__MODULE__) && !isdefined($mod, RGF._tagname) && RGF.init($mod)
        RuntimeGeneratedFunction($mod, $mod, expr)
    end
    ))
end

The only problem is that in small reproducer it works as expected, in Metatheory.jl tests it still failed. So, I had to add an Metatheory.init(mod::Module) function that calls closure_generator once. It feels quite hacky and I have a feeling that this could be greatly simplified.

Metatheory.init(mod) = closure_generator(mod)(:(x -> x))

Couldn't this behaviour be (sort of) integrated in the RGF package? My package tests (a lot, also about accessing values and dispatching methods in external modules in these "closures") are all passing after using this hack, the original issue was solved.

[ChrisRackauckas]

It probably could be, but I don't know if it needs to be.

[0x0f0f0f]

Recall the closures example from GeneralizedGenerated.jl :

@gg function h(x, c)
    quote
        d = x + 10
        function g(x, y=c)
            x + y + d
        end
    end
end

h(1, 2)(1) # => 14

I managed to get the same behaviour. ITS REALLY UGLY THO, and I'm not sure if this would hold at all if mechanized. I think it's worth a try.

cgen = closure_generator

h = cgen(@__MODULE__)(:( (x, c) -> begin
   d = x + 10
   cgen(@__MODULE__)( :((x) -> x + $c + $d) ) # this doesnt convince me
   end
))

println(h(1,2)(1)) # => 14 

[c42f]

Well... it's true you can get a certain kind of "closure" by creating an AST and interpolating the captures in there. Like any other use of RuntimeGeneratedFunction this will work in very particular circumstances where you don't need to create the AST more than a few times. But if you do this in a loop you'll quickly leak a lot of memory and may run into https://github.com/SciML/RuntimeGeneratedFunctions.jl/issues/13 or other problems.

To be honest, I feel like closure support within a RuntimeGeneratedFunction might be out of scope of this package as it requires the scope analysis pass of the compiler frontend. You can, however mix up normal closures and RGFs to bind slots of the RGF:

function foo(x)
    f = @RuntimeGeneratedFunction(:((y,z)->y+z^2))
    # bind `x` to the second slot of RGF `f`
    return y -> f(y,x)
end

When you're generating an AST, you've got control over whether that AST uses a closure as part of its implementation, so I think it may simply be best to avoid doing that if you want things to work smoothly with this package.

[c42f]

Could someone rename this issue to "Closures in RGFs" or something?

To capture some (lightly edited) comments I made in a slack discussion with @shashi -

if you've got existing tools for analyzing local variables within the RGF AST and figuring out which ones need to be closed over for inner functions, you could likely add a generic closure_args field to RuntimeGeneratedFunction and add some extra argument unpacking in the generated function. That would presumably be the way to go. It just seems like a lot of engineering effort, unless the "tools for analyzing local variables" exist already. ... My suggestion if you want to do this, is to look up some code from the ecosystem which already does this kind of variable analysis. I think the GeneralizedGenerated author had some such tools. IMO it's not worth reinventing (if the dependency isn't too heavy). Tracking closed-over variables is the kind of thing that seems simple, but in reality it's a rabbit hole of subtle language rules related to syntax desugaring.

[simeonschaub]

I am pretty sure OpaqueClosures can be used in generated functions, so perhaps that might be a better solution here?