diff --git a/base/reflection.jl b/base/reflection.jl index f654458446948..50e26e1cc040c 100644 --- a/base/reflection.jl +++ b/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 diff --git a/test/reflection.jl b/test/reflection.jl index b77e7d4e778a8..d9629147deef7 100644 --- a/test/reflection.jl +++ b/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)