Get rid of get_mi, eagerly infer CodeInstance for task calls

serenity4 · serenity4 · commit c1573d9d9046 · 2025-03-21T10:22:06.000-04:00
diff --git a/src/Cthulhu.jl b/src/Cthulhu.jl
@@ -507,10 +507,9 @@ function _descend(term::AbstractTerminal, interp::AbstractInterpreter, curs::Abs
     display_CI = true
     view_cmd = cthulhu_typed
     iostream = term.out_stream::IO
-    function additional_descend(new_mi::MethodInstance)
+    function additional_descend(new_ci::CodeInstance)
         new_interp = CthulhuInterpreter(interp)
-        do_typeinf!(new_interp, new_mi)
-        _descend(term, new_interp, new_mi;
+        _descend(term, new_interp, new_ci;
                  debuginfo, optimize, interruptexc, iswarn, hide_type_stable, remarks,
                  with_effects, exception_type,
                  inline_cost, type_annotations, annotate_source,
@@ -538,20 +537,21 @@ function _descend(term::AbstractTerminal, interp::AbstractInterpreter, curs::Abs
                         # but make something up that looks plausible.
                         callsites = Callsite[]
                         if display_CI
-                            exct = @static VERSION ≥ v"1.11.0-DEV.945" ? codeinst.exct : nothing
+                            exct = @static VERSION ≥ v"1.11.0-DEV.945" ? codeinst.exctype : nothing
                             callsite = Callsite(-1, EdgeCallInfo(codeinst, codeinst.rettype, get_effects(codeinst), exct), :invoke)
                             println(iostream)
                             println(iostream, "│ ─ $callsite")
                             println(iostream, "│  return ", Const(codeinst.rettype_const))
                             println(iostream)
                         end
+                        mi = codeinst.def
                         @goto show_menu
                     else
                         @info """
                         Inference discarded the source for this call because of recursion:
                         Cthulhu nevertheless is trying to retrieve the source for further inspection.
                         """
-                        additional_descend(get_mi(curs))
+                        additional_descend(get_ci(curs))
                         break
                     end
                 end
@@ -575,7 +575,7 @@ function _descend(term::AbstractTerminal, interp::AbstractInterpreter, curs::Abs
         end
         pc2excts = nothing
         end
-        callsites, sourcenodes = find_callsites(interp, src, infos, mi, slottypes, optimize & !annotate_source, annotate_source, pc2excts)
+        callsites, sourcenodes = find_callsites(interp, src, infos, ci, slottypes, optimize & !annotate_source, annotate_source, pc2excts)
 
         if jump_always
             if isdefined(Main, :VSCodeServer) && Main.VSCodeServer isa Module && isdefined(Main.VSCodeServer, :openfile)
@@ -671,7 +671,7 @@ function _descend(term::AbstractTerminal, interp::AbstractInterpreter, curs::Abs
                 if sourcenode !== nothing
                     show_sub_callsites = let callsite=callsite
                         map(info.callinfos) do ci
-                            p = Base.unwrap_unionall(get_mi(ci).specTypes).parameters
+                            p = Base.unwrap_unionall(ci.def.specTypes).parameters
                             if isa(sourcenode, TypedSyntax.MaybeTypedSyntaxNode) && length(p) == length(JuliaSyntax.children(sourcenode)) + 1
                                 newnode = copy(sourcenode)
                                 for (i, child) in enumerate(JuliaSyntax.children(newnode))
@@ -705,7 +705,7 @@ function _descend(term::AbstractTerminal, interp::AbstractInterpreter, curs::Abs
                 Inference didn't analyze this call because it is a dynamic call:
                 Cthulhu nevertheless is trying to descend into it for further inspection.
                 """
-                additional_descend(get_mi(info)::MethodInstance)
+                additional_descend(get_ci(info)::CodeInstance)
                 continue
             elseif info isa RTCallInfo
                 @info """
diff --git a/src/callsite.jl b/src/callsite.jl
@@ -2,26 +2,21 @@ using Unicode
 
 abstract type CallInfo end
 
-get_mi(ci::CallInfo) = get_ci(ci).def
-
 # Call could be resolved to a singular MI
 struct EdgeCallInfo <: CallInfo
-    ci::Union{CodeInstance,Nothing}
-    mi::MethodInstance
+    edge::CodeInstance
     rt
     effects::Effects
     exct
-    function EdgeCallInfo(edge::Union{MethodInstance, CodeInstance}, @nospecialize(rt), effects::Effects, @nospecialize(exct=nothing))
-        ci, mi = isa(edge, MethodInstance) ? (nothing, edge) : (edge, edge.def)
+    function EdgeCallInfo(edge::CodeInstance, @nospecialize(rt), effects::Effects, @nospecialize(exct=nothing))
         if isa(rt, LimitedAccuracy)
-            return LimitedCallInfo(new(ci, mi, ignorelimited(rt), effects, exct))
+            return LimitedCallInfo(new(edge, ignorelimited(rt), effects, exct))
         else
-            return new(ci, mi, rt, effects, exct)
+            return new(edge, rt, effects, exct)
         end
     end
 end
-get_ci(ci::EdgeCallInfo) = ci.ci
-get_mi(ci::EdgeCallInfo) = ci.mi
+get_ci(ci::EdgeCallInfo) = ci.edge
 get_rt(ci::EdgeCallInfo) = ci.rt
 get_effects(ci::EdgeCallInfo) = ci.effects
 get_exct(ci::EdgeCallInfo) = ci.exct
@@ -278,7 +273,7 @@ function Base.show(io::IO, (;exct)::ExctWrapper)
 end
 
 function show_callinfo(limiter, ci::EdgeCallInfo)
-    mi = ci.mi
+    mi = ci.edge.def
     tt = (Base.unwrap_unionall(mi.specTypes)::DataType).parameters[2:end]
     if !isa(mi.def, Method)
         name = ":toplevel"
@@ -460,12 +455,12 @@ _wrapped_callinfo(limiter, ::LimitedCallInfo)  = print(limiter, "limited")
 
 # is_callsite returns true if `call` dispatches to `callee`
 # See also `maybe_callsite` below
-is_callsite(call::MethodInstance, callee::MethodInstance) = call === callee
+is_callsite(call::CodeInstance, callee::MethodInstance) = call.def === callee
 is_callsite(::Nothing, callee::MethodInstance) = false   # for when `get_mi` returns `nothing`
 
 # is_callsite for higher-level inputs
 is_callsite(cs::Callsite, callee::MethodInstance) = is_callsite(cs.info, callee)
-is_callsite(info::CallInfo, callee::MethodInstance) = is_callsite(get_mi(info), callee)
+is_callsite(info::CallInfo, callee::MethodInstance) = is_callsite(get_ci(info), callee)
 # special CallInfo cases:
 function is_callsite(info::MultiCallInfo, callee::MethodInstance)
     for csi in info.callinfos
@@ -476,7 +471,7 @@ end
 
 # maybe_callsite returns true if `call` *might* dispatch to `callee`
 # See also `is_callsite` above
-function maybe_callsite(call::MethodInstance, callee::MethodInstance)
+function maybe_callsite(call::CodeInstance, callee::MethodInstance)
     # handle comparison among Varargs
     function generalized_va_subtype(@nospecialize(Tshort), @nospecialize(Tlong))
         nshort, nlong = length(Tshort.parameters), length(Tlong.parameters)
@@ -491,7 +486,7 @@ function maybe_callsite(call::MethodInstance, callee::MethodInstance)
         return T <: unwrapva(Tlong.parameters[end])
     end
 
-    Tcall, Tcallee = call.specTypes, callee.specTypes
+    Tcall, Tcallee = call.def.specTypes, callee.specTypes
     Tcall <: Tcallee && return true
     # Make sure we handle Tcall = Tuple{Vararg{String}}, Tcallee = Tuple{String,Vararg{String}}
     if Base.isvatuple(Tcall) && Base.isvatuple(Tcallee)
@@ -506,7 +501,7 @@ end
 # maybe_callsite for higher-level inputs
 maybe_callsite(cs::Callsite, callee::MethodInstance) = maybe_callsite(cs.info, callee)
 maybe_callsite(cs::Callsite, @nospecialize(tt::Type)) = maybe_callsite(cs.info, tt)
-maybe_callsite(info::CallInfo, callee::MethodInstance) = maybe_callsite(get_mi(info), callee)
+maybe_callsite(info::CallInfo, callee::MethodInstance) = maybe_callsite(get_ci(info), callee)
 # Special CallInfo cases:
 function maybe_callsite(info::MultiCallInfo, callee::MethodInstance)
     for csi in info.callinfos
diff --git a/src/codeview.jl b/src/codeview.jl
@@ -182,10 +182,10 @@ function cthulhu_typed(io::IO, debuginfo::Symbol,
             callsite_diagnostics = TypedSyntax.Diagnostic[]
             if (diagnostics_vscode || inlay_types_vscode)
                 vscode_io = IOContext(devnull, :inlay_hints=>vscode_io[:inlay_hints], :diagnostics=>vscode_io[:diagnostics])
-                callsite_mis = Dict() # type annotation is a bit long so I skipped it, doesn't seem to affect performance
-                visited_mis = Set{MethodInstance}((mi,))
-                add_callsites!(callsite_mis, visited_mis, callsite_diagnostics, mi; optimize, annotate_source, interp)
-                for callsite in values(callsite_mis)
+                callsite_cis = Dict() # type annotation is a bit long so I skipped it, doesn't seem to affect performance
+                visited_cis = Set{CodeInstance}((codeinst,))
+                add_callsites!(callsite_cis, visited_cis, callsite_diagnostics, codeinst; optimize, annotate_source, interp)
+                for callsite in values(callsite_cis)
                     if !isnothing(callsite)
                         descend_into_callsite!(vscode_io, callsite.tsn; iswarn, hide_type_stable, type_annotations)
                     end
@@ -333,13 +333,14 @@ function descend_into_callsite!(io::IO, tsn::TypedSyntaxNode;
     end
 end
 
-function add_callsites!(d::AbstractDict, visited_mis::AbstractSet, diagnostics::AbstractVector,
-    mi::MethodInstance, source_mi::MethodInstance=mi;
+function add_callsites!(d::AbstractDict, visited_cis::AbstractSet, diagnostics::AbstractVector,
+    ci::CodeInstance, source_ci::CodeInstance=ci;
     optimize::Bool=true, annotate_source::Bool=false,
     interp::AbstractInterpreter=CthulhuInterpreter())
+    mi = ci.def
 
     callsites, src, rt = try
-        (; src, rt, infos, slottypes, effects, codeinf) = lookup(interp, mi, optimize & !annotate_source)
+        (; src, rt, infos, slottypes, effects, codeinf) = lookup(interp, ci, optimize & !annotate_source)
 
         src = preprocess_ci!(src, mi, optimize & !annotate_source, CONFIG)
         if (optimize & !annotate_source) || isa(src, IRCode) # optimization might have deleted some statements
@@ -350,23 +351,25 @@ function add_callsites!(d::AbstractDict, visited_mis::AbstractSet, diagnostics::
 
         # We pass false as it doesn't affect callsites and skips fetching the method definition
         # using CodeTracking which is slow
-        callsites, _ = find_callsites(interp, src, infos, mi, slottypes, optimize & !annotate_source, false)
+        callsites, _ = find_callsites(interp, src, infos, ci, slottypes, optimize & !annotate_source, false)
         callsites, src, rt
     catch
         return nothing
     end
 
     for callsite in callsites
-        callsite_mi = callsite.info isa MultiCallInfo ? nothing : get_mi(callsite)
+        callsite_ci = callsite.info isa MultiCallInfo ? nothing : get_ci(callsite)
 
-        if !isnothing(callsite_mi) && callsite_mi ∉ visited_mis
-            push!(visited_mis, callsite_mi)
-            add_callsites!(d, visited_mis, diagnostics, callsite_mi, source_mi; optimize, annotate_source, interp)
+        if !isnothing(callsite_ci) && callsite_ci ∉ visited_cis
+            push!(visited_cis, callsite_ci)
+            # TODO: figure out why this `CodeInstance` is not present in the unoptimized cache.
+            callsite_ci.def.def.sig === Tuple{typeof(getproperty), Module, Symbol} && continue
+            add_callsites!(d, visited_cis, diagnostics, callsite_ci, source_ci; optimize, annotate_source, interp)
         end
     end
 
-    # Check if callsite is not just filling in default arguments and defined in same file as source_mi
-    if mi == source_mi || mi.def.file != source_mi.def.file
+    # Check if callsite is not just filling in default arguments and defined in same file as source_ci
+    if ci == source_ci || ci.def.def.file != source_ci.def.def.file
         return nothing
     end
     tsn, _ = get_typed_sourcetext(mi, src, rt; warn=false)
@@ -390,7 +393,7 @@ function add_callsites!(d::AbstractDict, visited_mis::AbstractSet, diagnostics::
             )
         end
     else
-        d[key] = (mi=mi, tsn=tsn)
+        d[key] = (ci=ci, tsn=tsn)
     end
 end
 
diff --git a/src/interface.jl b/src/interface.jl
@@ -74,8 +74,6 @@ get_pc_exct(interp::CthulhuInterpreter, key::InferenceKey) = get(interp.exceptio
 # a sensible default cursor for a MethodInstance
 AbstractCursor(interp::AbstractInterpreter, ci::CodeInstance) = CthulhuCursor(ci)
 
-get_mi(curs::AbstractCursor) = get_ci(curs).def
-
 mutable struct CustomToggle
     onoff::Bool
     key::UInt32
diff --git a/src/reflection.jl b/src/reflection.jl
@@ -6,18 +6,20 @@ using Base.Meta
 using Base: may_invoke_generator
 
 transform(::Val, callsite) = callsite
-function transform(::Val{:CuFunction}, callsite, callexpr, CI, mi, slottypes; world=get_world_counter())
+function transform(::Val{:CuFunction}, interp, callsite, callexpr, CI, mi, slottypes; world=get_world_counter())
     sptypes = sptypes_from_meth_instance(mi)
     tt = argextype(callexpr.args[4], CI, sptypes, slottypes)
     ft = argextype(callexpr.args[3], CI, sptypes, slottypes)
     isa(tt, Const) || return callsite
-    return Callsite(callsite.id, CuCallInfo(callinfo(Tuple{widenconst(ft), tt.val.parameters...}, Nothing; world)), callsite.head)
+    sig = Tuple{widenconst(ft), tt.val.parameters...}
+    return Callsite(callsite.id, CuCallInfo(callinfo(interp, sig, Nothing; world)), callsite.head)
 end
 
 function find_callsites(interp::AbstractInterpreter, CI::Union{CodeInfo,IRCode},
-                        stmt_infos::Union{Vector{CCCallInfo}, Nothing}, mi::MethodInstance,
+                        stmt_infos::Union{Vector{CCCallInfo}, Nothing}, ci::CodeInstance,
                         slottypes::Vector{Any}, optimize::Bool=true, annotate_source::Bool=false,
                         pc2excts::Union{Nothing,PC2Excts}=nothing)
+    mi = ci.def
     sptypes = sptypes_from_meth_instance(mi)
     callsites, sourcenodes = Callsite[], Union{TypedSyntax.MaybeTypedSyntaxNode,Callsite}[]
     if isa(CI, IRCode)
@@ -89,7 +91,7 @@ function find_callsites(interp::AbstractInterpreter, CI::Union{CodeInfo,IRCode},
                         func = args[6]
                         ftype = widenconst(argextype(func, CI, sptypes, slottypes))
                         sig = Tuple{ftype}
-                        callsite = Callsite(id, TaskCallInfo(callinfo(sig, nothing; world=get_inference_world(interp))), head)
+                        callsite = Callsite(id, TaskCallInfo(callinfo(interp, sig, nothing; world=get_inference_world(interp))), head)
                     end
                 end
             end
@@ -101,7 +103,7 @@ function find_callsites(interp::AbstractInterpreter, CI::Union{CodeInfo,IRCode},
                 ci = get_ci(info)
                 meth = ci.def.def
                 if isa(meth, Method) && nameof(meth.module) === :CUDAnative && meth.name === :cufunction
-                    callsite = transform(Val(:CuFunction), callsite, c, CI, ci.def, slottypes; world=get_inference_world(interp))
+                    callsite = transform(Val(:CuFunction), interp, callsite, c, CI, ci.def, slottypes; world=get_inference_world(interp))
                 end
             end
 
@@ -281,7 +283,7 @@ function preprocess_ci!(ir::IRCode, _::MethodInstance, optimize::Bool, config::C
     return ir
 end
 
-function callinfo(sig, rt, max_methods=-1; world=get_world_counter())
+function callinfo(interp, sig, rt, max_methods=-1; world=get_world_counter())
     methds = Base._methods_by_ftype(sig, max_methods, world)
     methds isa Bool && return FailedCallInfo(sig, rt)
     length(methds) < 1 && return FailedCallInfo(sig, rt)
@@ -295,7 +297,8 @@ function callinfo(sig, rt, max_methods=-1; world=get_world_counter())
         else
             mi = specialize_method(meth, atypes, sparams)
             if mi !== nothing
-                push!(callinfos, EdgeCallInfo(mi, rt, Effects()))
+                edge = do_typeinf!(interp, mi)
+                push!(callinfos, EdgeCallInfo(edge, rt, Effects()))
             else
                 push!(callinfos, FailedCallInfo(sig, rt))
             end
@@ -315,7 +318,7 @@ function find_caller_of(interp::AbstractInterpreter, callee::Union{MethodInstanc
     for optimize in (true, false)
         (; src, rt, infos, slottypes) = lookup(interp′, codeinst, optimize)
         src = preprocess_ci!(src, caller, optimize, CONFIG)
-        callsites, _ = find_callsites(interp′, src, infos, caller, slottypes, optimize)
+        callsites, _ = find_callsites(interp′, src, infos, codeinst, slottypes, optimize)
         callsites = allow_unspecialized ? filter(cs->maybe_callsite(cs, callee), callsites) :
                                           filter(cs->is_callsite(cs, callee), callsites)
         foreach(cs -> add_sourceline!(locs, src, cs.id, caller), callsites)
diff --git a/test/setup.jl b/test/setup.jl
@@ -22,11 +22,14 @@ end
 function find_callsites_by_ftt(@nospecialize(f), @nospecialize(TT=Tuple{}); optimize=true)
     (; interp, src, infos, codeinst, slottypes) = cthulhu_info(f, TT; optimize)
     src === nothing && return Cthulhu.Callsite[]
-    callsites, _ = Cthulhu.find_callsites(interp, src, infos, codeinst.def, slottypes, optimize)
+    callsites, _ = Cthulhu.find_callsites(interp, src, infos, codeinst, slottypes, optimize)
     @test all(c -> Cthulhu.get_effects(c) isa Cthulhu.Effects, callsites)
     return callsites
 end
 
 macro find_callsites_by_ftt(ex0...)
     return InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :find_callsites_by_ftt, ex0)
 end
+
+get_mi(callsite::Cthulhu.Callsite) = Cthulhu.get_ci(callsite).def
+get_method(callsite::Cthulhu.Callsite) = get_mi(callsite).def
diff --git a/test/test_Cthulhu.jl b/test/test_Cthulhu.jl
