add optional argument to code_lowered to enable generator expansion (#…

…22979)

base/reflection.jl

@@ -558,20 +558,43 @@ function to_tuple_type(@nospecialize(t))
     t
 end
 
-tt_cons(@nospecialize(t), @nospecialize(tup)) = (@_pure_meta; Tuple{t, (isa(tup, Type) ? tup.parameters : tup)...})
+function signature_type(@nospecialize(f), @nospecialize(args))
+    f_type = isa(f, Type) ? Type{f} : typeof(f)
+    arg_types = isa(args, Type) ? args.parameters : args
+    return Tuple{f_type, arg_types...}
+end
 
 """
-    code_lowered(f, types)
+    code_lowered(f, types, expand_generated = true)
+
+Return an array of lowered ASTs for the methods matching the given generic function and type signature.
+
+If `expand_generated` is `false`, then the `CodeInfo` instances returned for `@generated`
+methods will correspond to the generators' lowered ASTs. If `expand_generated` is `true`,
+these `CodeInfo` instances will correspond to the lowered ASTs of the method bodies yielded
+by expanding the generators.
 
-Returns an array of lowered ASTs for the methods matching the given generic function and type signature.
+Note that an error will be thrown if `types` are not leaf types when `expand_generated` is
+`true` and the corresponding method is a `@generated` method.
 """
-function code_lowered(@nospecialize(f), @nospecialize t = Tuple)
-    asts = map(methods(f, t)) do m
-        return uncompressed_ast(m::Method)
+function code_lowered(@nospecialize(f), @nospecialize(t = Tuple), expand_generated::Bool = true)
+    return map(method_instances(f, t)) do m
+        if expand_generated && isgenerated(m)
+            if isa(m, Core.MethodInstance)
+                return Core.Inference.get_staged(m)
+            else # isa(m, Method)
+                error("Could not expand generator for `@generated` method ", m, ". ",
+                      "This can happen if the provided argument types (", t, ") are ",
+                      "not leaf types, but the `expand_generated` argument is `true`.")
+            end
+        end
+        return uncompressed_ast(m)
     end
-    return asts
 end
 
+isgenerated(m::Method) = isdefined(m, :generator)
+isgenerated(m::Core.MethodInstance) = isgenerated(m.def)
+
 # low-level method lookup functions used by the compiler
 
 unionlen(x::Union) = unionlen(x.a) + unionlen(x.b)
@@ -582,8 +605,7 @@ _uniontypes(@nospecialize(x), ts) = (push!(ts, x); ts)
 uniontypes(@nospecialize(x)) = _uniontypes(x, Any[])
 
 function _methods(@nospecialize(f), @nospecialize(t), lim::Int, world::UInt)
-    ft = isa(f,Type) ? Type{f} : typeof(f)
-    tt = isa(t,Type) ? Tuple{ft, t.parameters...} : Tuple{ft, t...}
+    tt = signature_type(f, t)
     return _methods_by_ftype(tt, lim, world)
 end
 
@@ -635,8 +657,7 @@ end
 methods(f::Core.Builtin) = MethodList(Method[], typeof(f).name.mt)
 
 function methods_including_ambiguous(@nospecialize(f), @nospecialize(t))
-    ft = isa(f,Type) ? Type{f} : typeof(f)
-    tt = isa(t,Type) ? Tuple{ft, t.parameters...} : Tuple{ft, t...}
+    tt = signature_type(f, t)
     world = typemax(UInt)
     min = UInt[typemin(UInt)]
     max = UInt[typemax(UInt)]
@@ -686,9 +707,21 @@ function length(mt::MethodTable)
 end
 isempty(mt::MethodTable) = (mt.defs === nothing)
 
-uncompressed_ast(m::Method) = uncompressed_ast(m, isdefined(m,:source) ? m.source : m.generator.inferred)
+uncompressed_ast(m::Method) = uncompressed_ast(m, isdefined(m, :source) ? m.source : m.generator.inferred)
 uncompressed_ast(m::Method, s::CodeInfo) = s
 uncompressed_ast(m::Method, s::Array{UInt8,1}) = ccall(:jl_uncompress_ast, Any, (Any, Any), m, s)::CodeInfo
+uncompressed_ast(m::Core.MethodInstance) = uncompressed_ast(m.def)
+
+function method_instances(@nospecialize(f), @nospecialize(t), world::UInt = typemax(UInt))
+    tt = signature_type(f, t)
+    results = Vector{Union{Method,Core.MethodInstance}}()
+    for method_data in _methods_by_ftype(tt, -1, world)
+        mtypes, msp, m = method_data
+        instance = Core.Inference.code_for_method(m, mtypes, msp, world, false)
+        push!(results, ifelse(isa(instance, Core.MethodInstance), instance, m))
+    end
+    return results
+end
 
 # this type mirrors jl_cghooks_t (documented in julia.h)
 struct CodegenHooks
@@ -739,8 +772,7 @@ function _dump_function(@nospecialize(f), @nospecialize(t), native::Bool, wrappe
     world = typemax(UInt)
     meth = which(f, t)
     t = to_tuple_type(t)
-    ft = isa(f, Type) ? Type{f} : typeof(f)
-    tt = Tuple{ft, t.parameters...}
+    tt = signature_type(f, t)
     (ti, env) = ccall(:jl_type_intersection_with_env, Any, (Any, Any), tt, meth.sig)::SimpleVector
     meth = func_for_method_checked(meth, ti)
     linfo = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance}, (Any, Any, Any, UInt), meth, ti, env, world)
@@ -871,8 +903,7 @@ function which(@nospecialize(f), @nospecialize(t))
         length(ms)!=1 && error("no unique matching method for the specified argument types")
         return first(ms)
     else
-        ft = isa(f,Type) ? Type{f} : typeof(f)
-        tt = Tuple{ft, t.parameters...}
+        tt = signature_type(f, t)
         m = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, typemax(UInt))
         if m === nothing
             error("no method found for the specified argument types")
@@ -979,7 +1010,7 @@ true
 """
 function method_exists(@nospecialize(f), @nospecialize(t), world=typemax(UInt))
     t = to_tuple_type(t)
-    t = Tuple{isa(f,Type) ? Type{f} : typeof(f), t.parameters...}
+    t = signature_type(f, t)
     return ccall(:jl_method_exists, Cint, (Any, Any, UInt), typeof(f).name.mt, t, world) != 0
 end
 

test/reflection.jl

@@ -714,3 +714,32 @@ end
 @test_throws ErrorException fieldcount(Real)
 @test_throws ErrorException fieldcount(AbstractArray)
 @test_throws ErrorException fieldcount(Tuple{Any,Vararg{Any}})
+
+# PR #22979
+
+function test_similar_codeinfo(a, b)
+    @test a.code == b.code
+    @test a.slotnames == b.slotnames
+    @test a.slotflags == b.slotflags
+end
+
+@generated f22979(x...) = (y = 1; :(x[1] + x[2]))
+x22979 = (1, 2.0, 3.0 + im)
+T22979 = Tuple{typeof(f22979),typeof.(x22979)...}
+world = typemax(UInt)
+mtypes, msp, m = Base._methods_by_ftype(T22979, -1, world)[]
+instance = Core.Inference.code_for_method(m, mtypes, msp, world, false)
+cinfo_generated = Core.Inference.get_staged(instance)
+cinfo_ungenerated = Base.uncompressed_ast(m)
+
+test_similar_codeinfo(@code_lowered(f22979(x22979...)), cinfo_generated)
+
+cinfos = code_lowered(f22979, typeof.(x22979), true)
+@test length(cinfos) == 1
+cinfo = cinfos[]
+test_similar_codeinfo(cinfo, cinfo_generated)
+
+cinfos = code_lowered(f22979, typeof.(x22979), false)
+@test length(cinfos) == 1
+cinfo = cinfos[]
+test_similar_codeinfo(cinfo, cinfo_ungenerated)