stacktraces: Add an extension point for printing custom-owner CIs (#58430)

On Julia master, it is possible to add custom CodeInstances with non
`nothing` owner to the runtime system (e.g. by invoking them or with
custom AbstractInterpreters). However, the backtrace printing for these
custom code instances was not in any way distinguished from the
backtrace printing for an ordinary code instance of the same method
instance, making it sometimes hard to determine what code instance a
particular backtrace was referring to (if the downstream compiler was
making many different code instances for a particular MethodInstance).
Remidy this by adding an explicit extension point that can be
overwritten to modify the printing of these custom-owner CodeInstances.

Author: Keno

Compiler/src/Compiler.jl

@@ -65,7 +65,9 @@ using Base: @_foldable_meta, @_gc_preserve_begin, @_gc_preserve_end, @nospeciali
     structdiff, tls_world_age, unconstrain_vararg_length, unionlen, uniontype_layout,
     uniontypes, unsafe_convert, unwrap_unionall, unwrapva, vect, widen_diagonal,
     _uncompressed_ir, maybe_add_binding_backedge!, datatype_min_ninitialized,
-    partialstruct_init_undefs, fieldcount_noerror, _eval_import, _eval_using
+    partialstruct_init_undefs, fieldcount_noerror, _eval_import, _eval_using,
+    get_ci_mi
+
 using Base.Order
 
 import Base: ==, _topmod, append!, convert, copy, copy!, findall, first, get, get!,

Compiler/src/abstractinterpretation.jl

@@ -2221,12 +2221,6 @@ function get_ci_abi(ci::CodeInstance)
     (def::MethodInstance).specTypes
 end
 
-function get_ci_mi(ci::CodeInstance)
-    def = ci.def
-    isa(def, ABIOverride) && return def.def
-    return def::MethodInstance
-end
-
 function abstract_invoke(interp::AbstractInterpreter, arginfo::ArgInfo, si::StmtInfo, sv::AbsIntState)
     argtypes = arginfo.argtypes
     ft′ = argtype_by_index(argtypes, 2)

base/invalidation.jl

@@ -188,7 +188,7 @@ invalidate_code_for_globalref!(gr::GlobalRef, invalidated_bpart::Core.BindingPar
     invalidate_code_for_globalref!(convert(Core.Binding, gr), invalidated_bpart, new_bpart, new_max_world)
 
 function maybe_add_binding_backedge!(b::Core.Binding, edge::Union{Method, CodeInstance})
-    meth = isa(edge, Method) ? edge : Compiler.get_ci_mi(edge).def
+    meth = isa(edge, Method) ? edge : get_ci_mi(edge).def
     ccall(:jl_maybe_add_binding_backedge, Cint, (Any, Any, Any), b, edge, meth)
     return nothing
 end

base/runtime_internals.jl

@@ -1478,6 +1478,15 @@ end
 _uncompressed_ir(codeinst::CodeInstance, s::String) =
     ccall(:jl_uncompress_ir, Ref{CodeInfo}, (Any, Any, Any), codeinst.def.def::Method, codeinst, s)
 
+function get_ci_mi(codeinst::CodeInstance)
+    def = codeinst.def
+    if def isa Core.ABIOverride
+        return def.def
+    else
+        return def::MethodInstance
+    end
+end
+
 """
     Base.generating_output([incremental::Bool])::Bool
 

base/stacktraces.jl

@@ -273,8 +273,12 @@ function show_spec_linfo(io::IO, frame::StackFrame)
         if linfo isa Union{MethodInstance, CodeInstance}
             def = frame_method_or_module(frame)
             if def isa Module
-                Base.show_mi(io, linfo, #=from_stackframe=#true)
+                Base.show_mi(io, linfo::MethodInstance, #=from_stackframe=#true)
+            elseif linfo isa CodeInstance && linfo.owner !== nothing
+                show_custom_spec_sig(io, linfo.owner, linfo, frame)
             else
+                # Equivalent to the default implementation of `show_custom_spec_sig`
+                # for `linfo isa CodeInstance`, but saves an extra dynamic dispatch.
                 show_spec_sig(io, def, frame_mi(frame).specTypes)
             end
         else
@@ -284,6 +288,12 @@ function show_spec_linfo(io::IO, frame::StackFrame)
     end
 end
 
+# Can be extended by compiler packages to customize backtrace display of custom code instance frames
+function show_custom_spec_sig(io::IO, @nospecialize(owner), linfo::CodeInstance, frame::StackFrame)
+    mi = get_ci_mi(linfo)
+    return show_spec_sig(io, mi.def, mi.specTypes)
+end
+
 function show_spec_sig(io::IO, m::Method, @nospecialize(sig::Type))
     if get(io, :limit, :false)::Bool
         if !haskey(io, :displaysize)

base/staticdata.jl

@@ -3,22 +3,13 @@
 module StaticData
 
 using .Core: CodeInstance, MethodInstance
-using .Base: JLOptions, Compiler, get_world_counter, _methods_by_ftype, get_methodtable
+using .Base: JLOptions, Compiler, get_world_counter, _methods_by_ftype, get_methodtable, get_ci_mi
 
 const WORLD_AGE_REVALIDATION_SENTINEL::UInt = 1
 const _jl_debug_method_invalidation = Ref{Union{Nothing,Vector{Any}}}(nothing)
 debug_method_invalidation(onoff::Bool) =
     _jl_debug_method_invalidation[] = onoff ? Any[] : nothing
 
-function get_ci_mi(codeinst::CodeInstance)
-    def = codeinst.def
-    if def isa Core.ABIOverride
-        return def.def
-    else
-        return def::MethodInstance
-    end
-end
-
 # Restore backedges to external targets
 # `edges` = [caller1, ...], the list of worklist-owned code instances internally
 # `ext_ci_list` = [caller1, ...], the list of worklist-owned code instances externally