Support "functors" for code reflection utilities (#58891)

As a follow-up to #57911, this
updates:
 - `Base.method_instance`
 - `Base.method_instances`
 - `Base.code_ircode`
 - `Base.code_lowered`
 - `InteractiveUtils.code_llvm`
 - `InteractiveUtils.code_native`
 - `InteractiveUtils.code_warntype`
 
 to support "functor" invocations.
 
e.g. `code_llvm((Foo, Int, Int))` which corresponds to `(::Foo)(::Int,
::Int)`

Author: topolarity

Compiler/test/codegen.jl

@@ -22,7 +22,7 @@ end
 # The tests below assume a certain format and safepoint_on_entry=true breaks that.
 function get_llvm(@nospecialize(f), @nospecialize(t), raw=true, dump_module=false, optimize=true)
     params = Base.CodegenParams(safepoint_on_entry=false, gcstack_arg = false, debug_info_level=Cint(2))
-    d = InteractiveUtils._dump_function(f, t, false, false, raw, dump_module, :att, optimize, :none, false, params)
+    d = InteractiveUtils._dump_function(InteractiveUtils.ArgInfo(f, t), false, false, raw, dump_module, :att, optimize, :none, false, params)
     sprint(print, d)
 end
 

base/reflection.jl

@@ -16,7 +16,7 @@ The keyword `debuginfo` controls the amount of code metadata present in the outp
 Note that an error will be thrown if `types` are not concrete types when `generated` is
 `true` and any of the corresponding methods are an `@generated` method.
 """
-function code_lowered(@nospecialize(f), @nospecialize(t=Tuple); generated::Bool=true, debuginfo::Symbol=:default)
+function code_lowered(@nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}); generated::Bool=true, debuginfo::Symbol=:default)
     if @isdefined(IRShow)
         debuginfo = IRShow.debuginfo(debuginfo)
     elseif debuginfo === :default
@@ -28,7 +28,7 @@ function code_lowered(@nospecialize(f), @nospecialize(t=Tuple); generated::Bool=
     world = get_world_counter()
     world == typemax(UInt) && error("code reflection cannot be used from generated functions")
     ret = CodeInfo[]
-    for m in method_instances(f, t, world)
+    for m in method_instances(argtypes, world)
         if generated && hasgenerator(m)
             if may_invoke_generator(m)
                 code = ccall(:jl_code_for_staged, Ref{CodeInfo}, (Any, UInt, Ptr{Cvoid}), m, world, C_NULL)
@@ -46,12 +46,17 @@ function code_lowered(@nospecialize(f), @nospecialize(t=Tuple); generated::Bool=
     return ret
 end
 
+function code_lowered(@nospecialize(f), @nospecialize(t=Tuple); generated::Bool=true, debuginfo::Symbol=:default)
+    tt = signature_type(f, t)
+    return code_lowered(tt; generated, debuginfo)
+end
+
 # for backwards compat
 const uncompressed_ast = uncompressed_ir
 const _uncompressed_ast = _uncompressed_ir
 
-function method_instances(@nospecialize(f), @nospecialize(t), world::UInt)
-    tt = signature_type(f, t)
+function method_instances(@nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}), world::UInt)
+    tt = to_tuple_type(argtypes)
     results = Core.MethodInstance[]
     # this make a better error message than the typeassert that follows
     world == typemax(UInt) && error("code reflection cannot be used from generated functions")
@@ -62,15 +67,26 @@ function method_instances(@nospecialize(f), @nospecialize(t), world::UInt)
     return results
 end
 
-function method_instance(@nospecialize(f), @nospecialize(t);
-                         world=Base.get_world_counter(), method_table=nothing)
+function method_instances(@nospecialize(f), @nospecialize(t), world::UInt)
     tt = signature_type(f, t)
+    return method_instances(tt, world)
+end
+
+function method_instance(@nospecialize(argtypes::Union{Tuple,Type{<:Tuple}});
+                         world=Base.get_world_counter(), method_table=nothing)
+    tt = to_tuple_type(argtypes)
     mi = ccall(:jl_method_lookup_by_tt, Any,
                 (Any, Csize_t, Any),
                 tt, world, method_table)
     return mi::Union{Nothing, MethodInstance}
 end
 
+function method_instance(@nospecialize(f), @nospecialize(t);
+                         world=Base.get_world_counter(), method_table=nothing)
+    tt = signature_type(f, t)
+    return method_instance(tt; world, method_table)
+end
+
 default_debug_info_kind() = unsafe_load(cglobal(:jl_default_debug_info_kind, Cint))
 
 # this type mirrors jl_cgparams_t (documented in julia.h)
@@ -431,6 +447,11 @@ function code_ircode(@nospecialize(f), @nospecialize(types = default_tt(f)); kwa
     return code_ircode_by_type(tt; kwargs...)
 end
 
+function code_ircode(@nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}); kwargs...)
+    tt = to_tuple_type(argtypes)
+    return code_ircode_by_type(tt; kwargs...)
+end
+
 """
     code_ircode_by_type(types::Type{<:Tuple}; ...)
 

stdlib/InteractiveUtils/src/codeview.jl

@@ -20,6 +20,28 @@ const llstyle = Dict{Symbol, Tuple{Bool, Union{Symbol, Int}}}(
     :funcname    => (false, :light_yellow),
 )
 
+struct ArgInfo
+    oc::Union{Core.OpaqueClosure,Nothing}
+    tt::Type{<:Tuple}
+
+    # Construct from a function object + argtypes
+    function ArgInfo(@nospecialize(f), @nospecialize(t))
+        if isa(f, Core.Builtin)
+            throw(ArgumentError("argument is not a generic function"))
+        elseif f isa Core.OpaqueClosure
+            return new(f, Base.to_tuple_type(t))
+        else
+            return new(nothing, signature_type(f, t))
+        end
+    end
+
+    # Construct from argtypes (incl. arg0)
+    function ArgInfo(@nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}))
+        tt = Base.to_tuple_type(argtypes)
+        return new(nothing, tt)
+    end
+end
+
 function printstyled_ll(io::IO, x, s::Symbol, trailing_spaces="")
     printstyled(io, x, bold=llstyle[s][1], color=llstyle[s][2])
     print(io, trailing_spaces)
@@ -143,7 +165,7 @@ See the [`@code_warntype`](@ref man-code-warntype) section in the Performance Ti
 
 See also: [`@code_warntype`](@ref), [`code_typed`](@ref), [`code_lowered`](@ref), [`code_llvm`](@ref), [`code_native`](@ref).
 """
-function code_warntype(io::IO, @nospecialize(f), @nospecialize(tt=Base.default_tt(f));
+function code_warntype(io::IO, arginfo::ArgInfo;
                        world=Base.get_world_counter(),
                        interp::Base.Compiler.AbstractInterpreter=Base.Compiler.NativeInterpreter(world),
                        debuginfo::Symbol=:default, optimize::Bool=false, kwargs...)
@@ -152,13 +174,14 @@ function code_warntype(io::IO, @nospecialize(f), @nospecialize(tt=Base.default_t
     debuginfo = Base.IRShow.debuginfo(debuginfo)
     lineprinter = Base.IRShow.__debuginfo[debuginfo]
     nargs::Int = 0
-    if isa(f, Core.OpaqueClosure)
-        isa(f.source, Method) && (nargs = f.source.nargs)
-        print_warntype_codeinfo(io, Base.code_typed_opaque_closure(f, tt)[1]..., nargs;
+    if arginfo.oc !== nothing
+        (; oc, tt) = arginfo
+        isa(oc.source, Method) && (nargs = oc.source.nargs)
+        print_warntype_codeinfo(io, Base.code_typed_opaque_closure(oc, tt)[1]..., nargs;
                                 lineprinter, label_dynamic_calls = optimize)
         return nothing
     end
-    tt = Base.signature_type(f, tt)
+    tt = arginfo.tt
     matches = findall(tt, Base.Compiler.method_table(interp))
     matches === nothing && Base.raise_match_failure(:code_warntype, tt)
     for match in matches.matches
@@ -176,6 +199,8 @@ function code_warntype(io::IO, @nospecialize(f), @nospecialize(tt=Base.default_t
     end
     nothing
 end
+code_warntype(io::IO, @nospecialize(f), @nospecialize(tt=Base.default_tt(f)); kwargs...) = code_warntype(io, ArgInfo(f, tt); kwargs...)
+code_warntype(io::IO, @nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}); kwargs...) = code_warntype(io, ArgInfo(argtypes); kwargs...)
 code_warntype(args...; kwargs...) = (@nospecialize; code_warntype(stdout, args...; kwargs...))
 
 using Base: CodegenParams
@@ -189,33 +214,30 @@ const OC_MISMATCH_WARNING =
 
 # Printing code representations in IR and assembly
 
-function _dump_function(@nospecialize(f), @nospecialize(t), native::Bool, wrapper::Bool,
+function _dump_function(arginfo::ArgInfo, native::Bool, wrapper::Bool,
                         raw::Bool, dump_module::Bool, syntax::Symbol,
                         optimize::Bool, debuginfo::Symbol, binary::Bool,
                         params::CodegenParams=CodegenParams(debug_info_kind=Cint(0), debug_info_level=Cint(2), safepoint_on_entry=raw, gcstack_arg=raw))
     ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
-    if isa(f, Core.Builtin)
-        throw(ArgumentError("argument is not a generic function"))
-    end
     warning = ""
     # get the MethodInstance for the method match
-    if !isa(f, Core.OpaqueClosure)
+    if arginfo.oc === nothing
         world = Base.get_world_counter()
-        match = Base._which(signature_type(f, t); world)
+        match = Base._which(arginfo.tt; world)
         mi = Base.specialize_method(match)
         # TODO: use jl_is_cacheable_sig instead of isdispatchtuple
         isdispatchtuple(mi.specTypes) || (warning = GENERIC_SIG_WARNING)
     else
-        world = UInt64(f.world)
-        tt = Base.to_tuple_type(t)
-        if !isdefined(f.source, :source)
+        (; oc, tt) = arginfo
+        world = UInt64(oc.world)
+        if !isdefined(oc.source, :source)
             # OC was constructed from inferred source. There's only one
             # specialization and we can't infer anything more precise either.
-            world = f.source.primary_world
-            mi = f.source.specializations::Core.MethodInstance
-            Base.hasintersect(typeof(f).parameters[1], tt) || (warning = OC_MISMATCH_WARNING)
+            world = oc.source.primary_world
+            mi = oc.source.specializations::Core.MethodInstance
+            Base.hasintersect(typeof(oc).parameters[1], tt) || (warning = OC_MISMATCH_WARNING)
         else
-            mi = Base.specialize_method(f.source, Tuple{typeof(f.captures), tt.parameters...}, Core.svec())
+            mi = Base.specialize_method(oc.source, Tuple{typeof(oc.captures), tt.parameters...}, Core.svec())
             isdispatchtuple(mi.specTypes) || (warning = GENERIC_SIG_WARNING)
         end
     end
@@ -236,19 +258,19 @@ function _dump_function(@nospecialize(f), @nospecialize(t), native::Bool, wrappe
         end
         if isempty(str)
             # if that failed (or we want metadata), use LLVM to generate more accurate assembly output
-            if !isa(f, Core.OpaqueClosure)
+            if arginfo.oc === nothing
                 src = Base.Compiler.typeinf_code(Base.Compiler.NativeInterpreter(world), mi, true)
             else
-                src, rt = Base.get_oc_code_rt(nothing, f, tt, true)
+                src, rt = Base.get_oc_code_rt(nothing, arginfo.oc, arginfo.tt, true)
             end
             src isa Core.CodeInfo || error("failed to infer source for $mi")
             str = _dump_function_native_assembly(mi, src, wrapper, syntax, debuginfo, binary, raw, params)
         end
     else
-        if !isa(f, Core.OpaqueClosure)
+        if arginfo.oc === nothing
             src = Base.Compiler.typeinf_code(Base.Compiler.NativeInterpreter(world), mi, true)
         else
-            src, rt = Base.get_oc_code_rt(nothing, f, tt, true)
+            src, rt = Base.get_oc_code_rt(nothing, arginfo.oc, arginfo.tt, true)
         end
         src isa Core.CodeInfo || error("failed to infer source for $mi")
         str = _dump_function_llvm(mi, src, wrapper, !raw, dump_module, optimize, debuginfo, params)
@@ -311,16 +333,18 @@ Keyword argument `debuginfo` may be one of source (default) or none, to specify
 
 See also: [`@code_llvm`](@ref), [`code_warntype`](@ref), [`code_typed`](@ref), [`code_lowered`](@ref), [`code_native`](@ref).
 """
-function code_llvm(io::IO, @nospecialize(f), @nospecialize(types=Base.default_tt(f));
+function code_llvm(io::IO, arginfo::ArgInfo;
                    raw::Bool=false, dump_module::Bool=false, optimize::Bool=true, debuginfo::Symbol=:default,
                    params::CodegenParams=CodegenParams(debug_info_kind=Cint(0), debug_info_level=Cint(2), safepoint_on_entry=raw, gcstack_arg=raw))
-    d = _dump_function(f, types, false, false, raw, dump_module, :intel, optimize, debuginfo, false, params)
+    d = _dump_function(arginfo, false, false, raw, dump_module, :intel, optimize, debuginfo, false, params)
     if highlighting[:llvm] && get(io, :color, false)::Bool
         print_llvm(io, d)
     else
         print(io, d)
     end
 end
+code_llvm(io::IO, @nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}); kwargs...) = code_llvm(io, ArgInfo(argtypes); kwargs...)
+code_llvm(io::IO, @nospecialize(f), @nospecialize(types=Base.default_tt(f)); kwargs...) = code_llvm(io, ArgInfo(f, types); kwargs...)
 code_llvm(args...; kwargs...) = (@nospecialize; code_llvm(stdout, args...; kwargs...))
 
 """
@@ -337,17 +361,19 @@ generic function and type signature to `io`.
 
 See also: [`@code_native`](@ref), [`code_warntype`](@ref), [`code_typed`](@ref), [`code_lowered`](@ref), [`code_llvm`](@ref).
 """
-function code_native(io::IO, @nospecialize(f), @nospecialize(types=Base.default_tt(f));
+function code_native(io::IO, arginfo::ArgInfo;
                      dump_module::Bool=true, syntax::Symbol=:intel, raw::Bool=false,
                      debuginfo::Symbol=:default, binary::Bool=false,
                      params::CodegenParams=CodegenParams(debug_info_kind=Cint(0), debug_info_level=Cint(2), safepoint_on_entry=raw, gcstack_arg=raw))
-    d = _dump_function(f, types, true, false, raw, dump_module, syntax, true, debuginfo, binary, params)
+    d = _dump_function(arginfo, true, false, raw, dump_module, syntax, true, debuginfo, binary, params)
     if highlighting[:native] && get(io, :color, false)::Bool
         print_native(io, d)
     else
         print(io, d)
     end
 end
+code_native(io::IO, @nospecialize(argtypes::Union{Tuple,Type{<:Tuple}}); kwargs...) = code_native(io, ArgInfo(argtypes); kwargs...)
+code_native(io::IO, @nospecialize(f), @nospecialize(types=Base.default_tt(f)); kwargs...) = code_native(io, ArgInfo(f, types); kwargs...)
 code_native(args...; kwargs...) = (@nospecialize; code_native(stdout, args...; kwargs...))
 
 ## colorized IR and assembly printing

stdlib/InteractiveUtils/test/runtests.jl

@@ -592,7 +592,9 @@ end # module ReflectionTest
 # Issue #18883, code_llvm/code_native for generated functions
 @generated f18883() = nothing
 @test !isempty(sprint(code_llvm, f18883, Tuple{}))
+@test !isempty(sprint(code_llvm, (typeof(f18883),)))
 @test !isempty(sprint(code_native, f18883, Tuple{}))
+@test !isempty(sprint(code_native, (typeof(f18883),)))
 
 ix86 = r"i[356]86"
 
@@ -865,6 +867,27 @@ let # `default_tt` should work with any function with one method
     end); true)
 end
 
+let # specifying calls as argtypes (incl. arg0) should be supported
+    @test (code_warntype(devnull, (typeof(function ()
+        sin(42)
+    end),)); true)
+    @test (code_warntype(devnull, (typeof(function (a::Int)
+        sin(42)
+    end), Int)); true)
+    @test (code_llvm(devnull, (typeof(function ()
+        sin(42)
+    end),)); true)
+    @test (code_llvm(devnull, (typeof(function (a::Int)
+        sin(42)
+    end), Int)); true)
+    @test (code_native(devnull, (typeof(function ()
+        sin(42)
+    end),)); true)
+    @test (code_native(devnull, (typeof(function (a::Int)
+        sin(42)
+    end), Int)); true)
+end
+
 @testset "code_llvm on opaque_closure" begin
     let ci = code_typed(+, (Int, Int))[1][1]
         ir = Core.Compiler.inflate_ir(ci)

test/rebinding.jl

@@ -285,7 +285,7 @@ module RangeMerge
 
     function get_llvm(@nospecialize(f), @nospecialize(t), raw=true, dump_module=false, optimize=true)
         params = Base.CodegenParams(safepoint_on_entry=false, gcstack_arg = false, debug_info_level=Cint(2))
-        d = InteractiveUtils._dump_function(f, t, false, false, raw, dump_module, :att, optimize, :none, false, params)
+        d = InteractiveUtils._dump_function(InteractiveUtils.ArgInfo(f, t), false, false, raw, dump_module, :att, optimize, :none, false, params)
         sprint(print, d)
     end
 

test/reflection.jl

@@ -16,6 +16,11 @@ function test_ir_reflection(freflect, f, types)
     nothing
 end
 
+function test_ir_reflection(freflect, argtypes)
+    @test !isempty(freflect(argtypes))
+    nothing
+end
+
 function test_bin_reflection(freflect, f, types)
     iob = IOBuffer()
     freflect(iob, f, types)
@@ -27,6 +32,9 @@ end
 function test_code_reflection(freflect, f, types, tester)
     tester(freflect, f, types)
     tester(freflect, f, (types.parameters...,))
+    tt = Base.signature_type(f, types)
+    tester(freflect, tt)
+    tester(freflect, (tt.parameters...,))
     nothing
 end
 
@@ -43,6 +51,7 @@ end
 
 test_code_reflections(test_ir_reflection, code_lowered)
 test_code_reflections(test_ir_reflection, code_typed)
+test_code_reflections(test_ir_reflection, Base.code_ircode)
 
 io = IOBuffer()
 Base.print_statement_costs(io, map, (typeof(sqrt), Tuple{Int}))
@@ -682,6 +691,10 @@ end
 @test Base.code_typed_by_type(Tuple{Type{<:Val}})[2][2] == Val
 @test Base.code_typed_by_type(Tuple{typeof(sin), Float64})[1][2] === Float64
 
+# functor-like code_typed(...)
+@test Base.code_typed((Type{<:Val},))[2][2] == Val
+@test Base.code_typed((typeof(sin), Float64))[1][2] === Float64
+
 # New reflection methods in 0.6
 struct ReflectionExample{T<:AbstractFloat, N}
     x::Tuple{T, N}
@@ -1038,11 +1051,12 @@ _test_at_locals2(1,1,0.5f0)
 
 @testset "issue #31687" begin
     import InteractiveUtils._dump_function
+    import InteractiveUtils.ArgInfo
 
     @noinline f31687_child(i) = f31687_nonexistent(i)
     f31687_parent() = f31687_child(0)
     params = Base.CodegenParams()
-    _dump_function(f31687_parent, Tuple{},
+    _dump_function(ArgInfo(f31687_parent, Tuple{}),
                    #=native=#false, #=wrapper=#false, #=raw=#true,
                    #=dump_module=#true, #=syntax=#:att, #=optimize=#false, :none,
                    #=binary=#false)
@@ -1131,9 +1145,12 @@ end
     @test 1+1 == 2
     mi1 = Base.method_instance(+, (Int, Int))
     @test mi1.def.name == :+
-    # Note `jl_method_lookup` doesn't returns CNull if not found
-    mi2 = @ccall jl_method_lookup(Any[+, 1, 1]::Ptr{Any}, 3::Csize_t, Base.get_world_counter()::Csize_t)::Ref{Core.MethodInstance}
-    @test mi1 == mi2
+    mi2 = Base.method_instance((typeof(+), Int, Int))
+    @test mi2.def.name == :+
+    # Note `jl_method_lookup` doesn't return CNull if not found
+    mi3 = @ccall jl_method_lookup(Any[+, 1, 1]::Ptr{Any}, 3::Csize_t, Base.get_world_counter()::Csize_t)::Ref{Core.MethodInstance}
+    @test mi1 == mi3
+    @test mi2 == mi3
 end
 
 Base.@assume_effects :terminates_locally function issue41694(x::Int)