improve return_type handling (#240)

- improve `ReturnTypeCallInfo` printing (see the example below)
- don't error when trying to descend into a failed `return_type` site

> before
```julia
julia> only_ints(::Integer) = 1;

julia> descend(; optimize=false) do
           t1 = Base._return_type(only_ints, Tuple{Int})     # successful `return_type`
           t2 = Base._return_type(only_ints, Tuple{Float64}) # failed `return_type`
           t1, t2
       end
(::var"#3#4")() in Main at REPL[3]:2
Variables
  #self#::Core.Const(var"#3#4"())
  t2::Type{Union{}}
  t1::Type{Int64}

│ ─ %-1  = invoke #3()::Core.Const((Int64, Union{}))
    @ REPL[3]:2 within `#3`
1 ─ %1 = Base._return_type::Core.Const(Core.Compiler.return_type)
│   %2 = Core.apply_type(Main.Tuple, Main.Int)::Core.Const(Tuple{Int64})
│        (t1 = (%1)(Main.only_ints, %2))::Core.Const(Int64)
│   @ REPL[3]:3 within `#3`
│   %4 = Base._return_type::Core.Const(Core.Compiler.return_type)
│   %5 = Core.apply_type(Main.Tuple, Main.Float64)::Core.Const(Tuple{Float64})
│        (t2 = (%4)(Main.only_ints, %5))::Core.Const(Union{})
│   @ REPL[3]:4 within `#3`
│   %7 = Core.tuple(t1::Core.Const(Int64), t2::Core.Const(Union{}))::Core.Const((Int64, Union{}))
└──      return %7
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [o]ptimize, [w]arn, [h]ide type-stable statements, [d]ebuginfo, [r]emarks, [i]nlining costs, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
Advanced: dump [P]arams cache.
 • %3  = return_type < only_ints(::Int64)::Core.Const(Int64) >
   %6  = return_type < → return_type(::typeof(only_ints),::Type{Tuple{Float64}})::Core.Const(Union{}) >
   ↩
```

> this commit
```julia
julia> only_ints(::Integer) = 1;

julia> descend(; optimize=false) do
           t1 = Base._return_type(only_ints, Tuple{Int})     # successful `return_type`
           t2 = Base._return_type(only_ints, Tuple{Float64}) # failed `return_type`
           t1, t2
       end
(::var"#3#4")() in Main at REPL[3]:2
Variables
  #self#::Core.Const(var"#3#4"())
  t2::Type{Union{}}
  t1::Type{Int64}

│ ─ %-1  = invoke #3()::Core.Const((Int64, Union{}))
    @ REPL[3]:2 within `#3`
1 ─ %1 = Base._return_type::Core.Const(Core.Compiler.return_type)
│   %2 = Core.apply_type(Main.Tuple, Main.Int)::Core.Const(Tuple{Int64})
│        (t1 = (%1)(Main.only_ints, %2))::Core.Const(Int64)
│   @ REPL[3]:3 within `#3`
│   %4 = Base._return_type::Core.Const(Core.Compiler.return_type)
│   %5 = Core.apply_type(Main.Tuple, Main.Float64)::Core.Const(Tuple{Float64})
│        (t2 = (%4)(Main.only_ints, %5))::Core.Const(Union{})
│   @ REPL[3]:4 within `#3`
│   %7 = Core.tuple(t1::Core.Const(Int64), t2::Core.Const(Union{}))::Core.Const((Int64, Union{}))
└──      return %7
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [o]ptimize, [w]arn, [h]ide type-stable statements, [d]ebuginfo, [r]emarks, [i]nlining costs, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
Advanced: dump [P]arams cache.
 • %3  = return_type < only_ints(::Int64)::Int64 >
   %6  = return_type < only_ints(::Float64)::Union{} >
   ↩
```

Author: aviatesk

src/Cthulhu.jl

@@ -9,7 +9,7 @@ using REPL: REPL, AbstractTerminal
 
 using Core: MethodInstance
 const Compiler = Core.Compiler
-using Core.Compiler: MethodMatch, LimitedAccuracy, ignorelimited
+import Core.Compiler: MethodMatch, LimitedAccuracy, ignorelimited, specialize_method
 import Base: unwrapva, isvarargtype
 const mapany = Base.mapany
 
@@ -462,7 +462,7 @@ end
 
 function get_specialization(@nospecialize(TT))
     match = Base._which(TT)
-    mi = Core.Compiler.specialize_method(match)
+    mi = specialize_method(match)
     return mi
 end
 

src/callsite.jl

@@ -1,6 +1,6 @@
 using Unicode
 
-abstract type CallInfo; end
+abstract type CallInfo end
 
 # Call could be resolved to a singular MI
 struct MICallInfo <: CallInfo
@@ -97,10 +97,10 @@ get_rt(tci::OCCallInfo) = get_rt(tci.ci)
 
 # Special handling for ReturnTypeCall
 struct ReturnTypeCallInfo <: CallInfo
-    called_mi::CallInfo
+    vmi::CallInfo # virtualized method call
 end
-get_mi(rtci::ReturnTypeCallInfo) = get_mi(rtci.called_mi)
-get_rt(rtci::ReturnTypeCallInfo) = get_rt(rtci.called_mi)
+get_mi((; vmi)::ReturnTypeCallInfo) = isa(vmi, FailedCallInfo) ? nothing : get_mi(vmi)
+get_rt((; vmi)::ReturnTypeCallInfo) = Type{isa(vmi, FailedCallInfo) ? Union{} : widenconst(get_rt(vmi))}
 
 struct ConstPropCallInfo <: CallInfo
     mi::CallInfo
@@ -232,7 +232,7 @@ function show_callinfo(limiter, mici::MICallInfo)
     mi = mici.mi
     tt = (Base.unwrap_unionall(mi.specTypes)::DataType).parameters[2:end]
     name = (mi.def::Method).name
-    rt = mici.rt
+    rt = get_rt(mici)
     __show_limited(limiter, name, tt, rt)
 end
 
@@ -247,8 +247,7 @@ function show_callinfo(limiter, ci::Union{MultiCallInfo, FailedCallInfo, Generat
     else
         name = "→ (::$(nameof(ft)))"
     end
-    rt = ci.rt
-    __show_limited(limiter, name::String, tt, rt)
+    __show_limited(limiter, name::String, tt, get_rt(ci))
 end
 
 function show_callinfo(limiter, (; argtypes, rt)::PureCallInfo)
@@ -262,7 +261,19 @@ function show_callinfo(limiter, ci::ConstPropCallInfo)
     # XXX: The first argument could be const-overriden too
     name = ci.result.linfo.def.name
     tt = ci.result.argtypes[2:end]
-    __show_limited(limiter, name, tt, (ignorewrappers(ci.mi)::MICallInfo).rt)
+    __show_limited(limiter, name, tt, get_rt(ignorewrappers(ci.mi)::MICallInfo))
+end
+
+function show_callinfo(limiter, (; vmi)::ReturnTypeCallInfo)
+    if isa(vmi, FailedCallInfo)
+        ft = Base.tuple_type_head(vmi.sig)
+        f = Compiler.singleton_type(ft)
+        name = isnothing(f) ? "unknown" : string(f)
+        tt = Base.tuple_type_tail(vmi.sig).parameters
+        __show_limited(limiter, name, tt, vmi.rt)
+    else
+        show_callinfo(limiter, vmi)
+    end
 end
 
 function Base.show(io::IO, c::Callsite)
@@ -306,7 +317,7 @@ function Base.show(io::IO, c::Callsite)
         print(limiter, " >")
     elseif isa(info, ReturnTypeCallInfo)
         print(limiter, " = return_type < ")
-        show_callinfo(limiter, info.called_mi)
+        show_callinfo(limiter, info)
         print(limiter, " >")
     elseif isa(info, CuCallInfo)
         print(limiter, " = cucall < ")
@@ -342,7 +353,7 @@ is_callsite(info::WrappedCallInfo, mi::MethodInstance) = is_callsite(get_wrapped
 is_callsite(info::ConstPropCallInfo, mi::MethodInstance) = is_callsite(info.mi, mi)
 is_callsite(info::TaskCallInfo, mi::MethodInstance) = is_callsite(info.ci, mi)
 is_callsite(info::InvokeCallInfo, mi::MethodInstance) = is_callsite(info.ci, mi)
-is_callsite(info::ReturnTypeCallInfo, mi::MethodInstance) = is_callsite(info.called_mi, mi)
+is_callsite(info::ReturnTypeCallInfo, mi::MethodInstance) = is_callsite(info.vmi, mi)
 is_callsite(info::CuCallInfo, mi::MethodInstance) = is_callsite(info.cumi, mi)
 function is_callsite(info::MultiCallInfo, mi::MethodInstance)
     for csi in info.callinfos

src/interpreter.jl

@@ -76,7 +76,7 @@ function Compiler.finish(state::InferenceState, interp::CthulhuInterpreter)
 end
 
 function Compiler.transform_result_for_cache(interp::CthulhuInterpreter, linfo::MethodInstance,
-        valid_worlds::Core.Compiler.WorldRange, @nospecialize(inferred_result))
+        valid_worlds::WorldRange, @nospecialize(inferred_result))
     if isa(inferred_result, OptimizationState)
         opt = inferred_result
         if isdefined(opt, :ir)

src/reflection.jl

@@ -5,16 +5,16 @@
 using Base.Meta
 import .Compiler: widenconst, argextype, Const, MethodMatchInfo,
     UnionSplitApplyCallInfo, UnionSplitInfo, ConstCallInfo,
-    MethodResultPure, ApplyCallInfo
+    MethodResultPure, ApplyCallInfo,
+    sptypes_from_meth_instance, argtypes_to_type
+import Base: may_invoke_generator
 
-const sptypes_from_meth_instance = Core.Compiler.sptypes_from_meth_instance
-const may_invoke_generator = Base.may_invoke_generator
 function code_for_method(method, metharg, methsp, world, preexisting=false)
     @static if VERSION ≥ v"1.8.0-DEV.369"
         # https://github.com/JuliaLang/julia/pull/41920
-        Core.Compiler.specialize_method(method, metharg, methsp; preexisting)
+        specialize_method(method, metharg, methsp; preexisting)
     else
-        Core.Compiler.specialize_method(method, metharg, methsp, preexisting)
+        specialize_method(method, metharg, methsp, preexisting)
     end
 end
 
@@ -27,6 +27,8 @@ function transform(::Val{:CuFunction}, callsite, callexpr, CI, mi, slottypes; pa
     return Callsite(callsite.id, CuCallInfo(callinfo(Tuple{widenconst(ft), tt.val.parameters...}, Nothing, params)), callsite.head)
 end
 
+const ArgTypes = Vector{Any}
+
 function find_callsites(interp::CthulhuInterpreter, CI::Union{Core.CodeInfo, IRCode},
                         stmt_info::Union{Vector, Nothing}, mi::Core.MethodInstance,
                         slottypes::Vector{Any}, optimize::Bool=true;
@@ -47,28 +49,17 @@ function find_callsites(interp::CthulhuInterpreter, CI::Union{Core.CodeInfo, IRC
                 if !optimize
                     args = (ignorelhs(c)::Expr).args
                 end
-                types = mapany(function (@nospecialize(arg),)
-                                   t = argextype(arg, CI, sptypes, slottypes)
-                                   return widenconst(ignorelimited(t))
-                               end, args)
-                was_return_type = false
-                if isa(info, Core.Compiler.ReturnTypeCallInfo)
-                    info = info.info
-                    was_return_type = true
-                end
-
-                callinfos = process_info(interp, info, types, rt, optimize)
-                if !was_return_type && isempty(callinfos)
-                    continue
-                end
+                argtypes = mapany(function (@nospecialize(arg),)
+                                      t = argextype(arg, CI, sptypes, slottypes)
+                                      return widenconst(ignorelimited(t))
+                                  end, args)
+                callinfos = process_info(interp, info, argtypes, rt, optimize)
+                isempty(callinfos) && continue
                 callsite = let
                     if length(callinfos) == 1
                         callinfo = callinfos[1]
                     else
-                        callinfo = MultiCallInfo(Compiler.argtypes_to_type(types), rt, callinfos)
-                    end
-                    if was_return_type
-                        callinfo = ReturnTypeCallInfo(callinfo)
+                        callinfo = MultiCallInfo(argtypes_to_type(argtypes), rt, callinfos)
                     end
                     Callsite(id, callinfo, c.head)
                 end
@@ -113,10 +104,9 @@ function find_callsites(interp::CthulhuInterpreter, CI::Union{Core.CodeInfo, IRC
     return callsites
 end
 
-function process_info(interp, @nospecialize(info), types, @nospecialize(rt), optimize::Bool)
+function process_info(interp, @nospecialize(info), argtypes::ArgTypes, @nospecialize(rt), optimize::Bool)
     is_cached(@nospecialize(key)) = haskey(optimize ? interp.opt : interp.unopt, key)
-    process_recursive(@nospecialize(newinfo)) = process_info(interp, newinfo, types, rt, optimize)
-
+    process_recursive(@nospecialize(newinfo)) = process_info(interp, newinfo, argtypes, rt, optimize)
 
     if isa(info, MethodMatchInfo)
         if info.results === missing
@@ -125,12 +115,12 @@ function process_info(interp, @nospecialize(info), types, @nospecialize(rt), opt
 
         matches = info.results.matches
         return mapany(matches) do match::Core.MethodMatch
-            mi = Core.Compiler.specialize_method(match)
+            mi = specialize_method(match)
             mici = MICallInfo(mi, rt)
             return is_cached(mi) ? mici : UncachedCallInfo(mici)
         end
     elseif isa(info, MethodResultPure)
-        return Any[PureCallInfo(types, rt)]
+        return Any[PureCallInfo(argtypes, rt)]
     elseif isa(info, UnionSplitInfo)
         return mapreduce(process_recursive, vcat, info.matches; init=[])::Vector{Any}
     elseif isa(info, UnionSplitApplyCallInfo)
@@ -158,16 +148,33 @@ function process_info(interp, @nospecialize(info), types, @nospecialize(rt), opt
     elseif (@static isdefined(Compiler, :OpaqueClosureCreateInfo) && true) && isa(info, Compiler.OpaqueClosureCreateInfo)
         # TODO: Add ability to descend into OCs at creation site
         return []
+    elseif isa(info, Compiler.ReturnTypeCallInfo)
+        newargtypes = argtypes[2:end]
+        callinfos = process_info(interp, info.info, newargtypes, unwrapType(widenconst(rt)), optimize)
+        if length(callinfos) == 1
+            vmi = only(callinfos)
+        else
+            @assert isempty(callinfos)
+            argt = unwrapType(widenconst(newargtypes[2]))::DataType
+            sig = Tuple{widenconst(newargtypes[1]), argt.parameters...}
+            vmi = FailedCallInfo(sig, Union{})
+        end
+        return Any[ReturnTypeCallInfo(vmi)]
     elseif info == false
         return []
     else
-        @show CI
-        @show c
-        @show info
-        error()
+        @eval Main begin
+            interp = $interp
+            info = $info
+            argtypes = $argtypes
+            rt = $rt
+            optimize = $optimize
+        end
+        error("inspect `Main.interp|info|argtypes|rt|optimize`")
     end
 end
 
+unwrapType(@nospecialize t) = Compiler.isType(t) ? t.parameters[1] : t
 
 ignorelhs(@nospecialize(x)) = isexpr(x, :(=)) ? last(x.args) : x
 function is_call_expr(x::Expr, optimize::Bool)

test/runtests.jl

@@ -234,22 +234,34 @@ struct SingletonPureCallable{N} end
     @test s == "SingletonPureCallable{1}()(::Float64)::Float64"
 end
 
-# Failed return_type
-only_ints(::Integer) = 1
-return_type_failure(::T) where T = Base._return_type(only_ints, Tuple{T})
-let callsites = find_callsites_by_ftt(return_type_failure, Tuple{Float64}, optimize=false)
-    @test length(callsites) == 1
-    callinfo = callsites[1].info
-    @test callinfo isa Cthulhu.ReturnTypeCallInfo
-    callinfo = callinfo.called_mi
-    @test callinfo isa Cthulhu.MultiCallInfo
+@testset "ReturnTypeCallInfo" begin
+    only_ints(::Integer) = 1
+
+    callsites = find_callsites_by_ftt(; optimize=false) do
+            t1 = Base._return_type(only_ints, Tuple{Int})     # successful `return_type`
+            t2 = Base._return_type(only_ints, Tuple{Float64}) # failed `return_type`
+            t1, t2
+        end
+    @test length(callsites) == 2
+    callinfo1 = callsites[1].info
+    @test callinfo1 isa Cthulhu.ReturnTypeCallInfo
+    @test callinfo1.vmi isa Cthulhu.MICallInfo
     io = IOBuffer()
-    Cthulhu.show_callinfo(io, callinfo)
-    @test String(take!(io)) == "→ return_type(::typeof(only_ints),::Type{Tuple{Float64}})::Core.Const(Union{})"
+    Cthulhu.show_callinfo(io, callinfo1)
+    @test String(take!(io)) == "only_ints(::$Int)::$Int"
     io = IOBuffer()
     print(io, callsites[1])
-    @test occursin("return_type < → return_type", String(take!(io)))
-    @test length(callinfo.callinfos) == 0
+    @test occursin("return_type < only_ints(::$Int)::$Int >", String(take!(io)))
+
+    callinfo2 = callsites[2].info
+    @test callinfo2 isa Cthulhu.ReturnTypeCallInfo
+    @test callinfo2.vmi isa Cthulhu.FailedCallInfo
+    io = IOBuffer()
+    Cthulhu.show_callinfo(io, callinfo2)
+    @test String(take!(io)) == "only_ints(::Float64)::Union{}"
+    io = IOBuffer()
+    print(io, callsites[2])
+    @test occursin("return_type < only_ints(::Float64)::Union{} >", String(take!(io)))
 end
 
 # tasks