Actually step into generators for generated functions

Author: timholy

src/construct.jl

@@ -47,11 +47,13 @@ end
 
 get_source(meth::Method) = Base.uncompressed_ast(meth)
 
-function get_source(g::GeneratedFunctionStub, env)
-    b = g(env..., g.argnames...)
-    b isa CodeInfo && return b
-    return eval(b)
-end
+# The following creates the CodeInfo holding the generated expression. May be
+# useful if we decide to revert to that behavior.
+# function get_source(g::GeneratedFunctionStub, env)
+#     b = @which g(env..., g.argnames...)
+#     b isa CodeInfo && return b
+#     return eval(b)
+# end
 
 function copy_codeinfo(code::CodeInfo)
     @static if VERSION < v"1.1.0-DEV.762"
@@ -129,6 +131,7 @@ function prepare_framecode(method::Method, @nospecialize(argtypes); enter_genera
         framecode = get(framedict, method, nothing)
     end
     if framecode === nothing
+        method0 = method
         if is_generated(method) && !enter_generated
             # If we're stepping into a staged function, we need to use
             # the specialization, rather than stepping through the
@@ -138,16 +141,20 @@ function prepare_framecode(method::Method, @nospecialize(argtypes); enter_genera
             generator = false
         else
             if is_generated(method)
-                code = get_source(method.generator, lenv)
+                g = method.generator
+                methsg = collect(methods(g.gen))
+                @assert length(methsg) == 1
+                method = first(methsg)
+                code = get_source(method)
                 generator = true
             else
                 code = get_source(method)
                 generator = false
             end
         end
         framecode = FrameCode(method, code; generator=generator)
-        if is_generated(method) && !enter_generated
-            genframedict[(method, argtypes)] = framecode
+        if is_generated(method0) && !generator
+            genframedict[(method0, argtypes)] = framecode
         else
             framedict[method] = framecode
         end
@@ -225,8 +232,10 @@ function prepare_call(@nospecialize(f), allargs; enter_generated = false)
     isa(ret, Compiled) && return ret
     # Typical return
     framecode, lenv = ret
-    if is_generated(method) && enter_generated
+    if framecode.generator
         args = Any[_Typeof(a) for a in args]
+        selfarg = Base.unwrap_unionall(scopeof(framecode).sig).parameters[1]  #  #self#
+        args = Any[selfarg, lenv..., args...]
     end
     return framecode, args, lenv, argtypes
 end
@@ -294,8 +303,9 @@ static parameters `lenv`. See [`JuliaInterpreter.prepare_call`](@ref) for inform
 """
 function prepare_frame(framecode::FrameCode, args::Vector{Any}, lenv::SimpleVector; enter_generated=false)
     s = scopeof(framecode)
-    if isa(s, Method) && is_generated(s) && enter_generated
+    if framecode.generator
         args = Any[_Typeof(a) for a in args]
+        args = Any[Base.unwrap_unionall(s.sig).parameters[1], lenv..., args...]  # first is #self#
     end
     framedata = prepare_framedata(framecode, args)
     resize!(framedata.sparams, length(lenv))

test/debug.jl

@@ -29,10 +29,14 @@ step_through(f, args...; kwargs...) = step_through_frame(() -> enter_call(f, arg
 step_through(expr::Expr) = step_through_frame(() -> enter_call_expr(expr))
 
 @generated function generatedfoo(x)
-    :(return x)
+    # A deliberately-complicated way to perform this operation
+    ex = Expr(:block)
+    push!(ex.args, :(return x))
+    return ex
 end
 callgenerated() = generatedfoo(1)
 @generated function generatedparams(a::Array{T,N}) where {T,N}
+    zz = 1
     :(return ($T,$N))
 end
 callgeneratedparams() = generatedparams([1 2; 3 4])
@@ -99,26 +103,37 @@ struct B{T} end
         frame = enter_call_expr(:($(callgenerated)()))
         f, pc = debug_command(frame, :sg)
         @test isa(pc, BreakpointRef)
-        @test JuliaInterpreter.scopeof(f).name == :generatedfoo
         stmt = JuliaInterpreter.pc_expr(f)
-        @test stmt.head == :return
+        if stmt === nothing
+            f, pc = debug_command(f, :se)
+            stmt = JuliaInterpreter.pc_expr(f)
+        end
+        @test stmt.head == :(=)
         f2, pc = debug_command(f, :finish)
         @test JuliaInterpreter.scopeof(f2).name == :callgenerated
         # Now finish the regular function
         @test debug_command(frame, :finish) === nothing
         @test frame.callee === nothing
-        @test get_return(frame) === Int
+        excmp = quote return x end
+        deleteat!(excmp.args, 1)  # delete LineNumberNode
+        @test get_return(frame) == excmp
 
         # Parametric generated function (see #157)
         frame = fr = JuliaInterpreter.enter_call(callgeneratedparams)
         while fr.pc < JuliaInterpreter.nstatements(fr.framecode) - 1
             fr, pc = debug_command(fr, :se)
         end
         fr, pc = debug_command(fr, :sg)
-        @test JuliaInterpreter.scopeof(fr).name == :generatedparams
+        JuliaInterpreter.finish!(fr)
+        vzz = filter(v -> v.name == :zz, JuliaInterpreter.locals(fr))[1]
+        @test vzz.value == 1
         fr, pc = debug_command(fr, :finish)
         @test debug_command(fr, :finish) === nothing
-        @test JuliaInterpreter.get_return(fr) == (Int, 2)
+        ex = get_return(fr)
+        isa(ex.args[1], LineNumberNode) && deleteat!(ex.args, 1)
+        excmp = quote return ($Int, 2) end
+        deleteat!(excmp.args, 1)
+        @test ex == excmp
     end
 
     @testset "Optional arguments" begin