add versions of code_typed and which that accept tuple types (#36389

)

base/reflection.jl

@@ -1081,10 +1081,31 @@ function code_typed(@nospecialize(f), @nospecialize(types=Tuple);
                     debuginfo::Symbol=:default,
                     world = get_world_counter(),
                     interp = Core.Compiler.NativeInterpreter(world))
-    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
+    ft = Core.Typeof(f)
+    if isa(types, Type)
+        u = unwrap_unionall(types)
+        tt = rewrap_unionall(Tuple{ft, u.parameters...}, types)
+    else
+        tt = Tuple{ft, types...}
+    end
+    return code_typed_by_type(tt; optimize, debuginfo, world, interp)
+end
+
+"""
+    code_typed_by_type(types::Type{<:Tuple}; ...)
+
+Similar to [`code_typed`](@ref), except the argument is a tuple type describing
+a full signature to query.
+"""
+function code_typed_by_type(@nospecialize(tt::Type);
+                            optimize=true,
+                            debuginfo::Symbol=:default,
+                            world = get_world_counter(),
+                            interp = Core.Compiler.NativeInterpreter(world))
+    ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
     if @isdefined(IRShow)
         debuginfo = IRShow.debuginfo(debuginfo)
     elseif debuginfo === :default
@@ -1093,10 +1114,14 @@ function code_typed(@nospecialize(f), @nospecialize(types=Tuple);
     if debuginfo !== :source && debuginfo !== :none
         throw(ArgumentError("'debuginfo' must be either :source or :none"))
     end
-    types = to_tuple_type(types)
+    tt = to_tuple_type(tt)
+    meths = _methods_by_ftype(tt, -1, world)
+    if meths === false
+        error("signature does not correspond to a generic function")
+    end
     asts = []
-    for x in _methods(f, types, -1, world)
-        meth = func_for_method_checked(x[3], types, x[2])
+    for x in meths
+        meth = func_for_method_checked(x[3], tt, x[2])
         (code, ty) = Core.Compiler.typeinf_code(interp, meth, x[1], x[2], optimize)
         code === nothing && error("inference not successful") # inference disabled?
         debuginfo === :none && remove_linenums!(code)
@@ -1135,6 +1160,15 @@ function which(@nospecialize(f), @nospecialize(t))
     end
     t = to_tuple_type(t)
     tt = signature_type(f, t)
+    return which(tt)
+end
+
+"""
+    which(types::Type{<:Tuple})
+
+Returns the method that would be called by the given type signature (as a tuple type).
+"""
+function which(@nospecialize(tt::Type))
     m = ccall(:jl_gf_invoke_lookup, Any, (Any, UInt), tt, typemax(UInt))
     if m === nothing
         error("no unique matching method found for the specified argument types")

test/reflection.jl

@@ -525,6 +525,11 @@ let
     @test !isdefined(mi.cache, :next)
 end
 
+# code_typed_by_type
+@test Base.code_typed_by_type(Tuple{Type{<:Val}})[1][2] == Val
+@test Base.code_typed_by_type(Tuple{typeof(sin), Float64})[1][2] === Float64
+@test_throws ErrorException("signature does not correspond to a generic function") Base.code_typed_by_type(Tuple{Any})
+
 # New reflection methods in 0.6
 struct ReflectionExample{T<:AbstractFloat, N}
     x::Tuple{T, N}
@@ -820,6 +825,7 @@ f20872(::Val, ::Val) = false
 @test which(f20872, Tuple{Val,Val}).sig == Tuple{typeof(f20872), Val, Val}
 @test which(f20872, Tuple{Val,Val{N}} where N).sig == Tuple{typeof(f20872), Val, Val}
 @test_throws ErrorException which(f20872, Tuple{Any,Val{N}} where N)
+@test which(Tuple{typeof(f20872), Val{1}, Val{2}}).sig == Tuple{typeof(f20872), Val, Val}
 
 module M29962 end
 # make sure checking if a binding is deprecated does not resolve it