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

…uliaLang#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)