invoked calls: record invoke signature in backedges (#46010)

This fixes a long-standing issue with how we've handled `invoke` calls
with respect to method invalidation.  When we load a package, we need
to ask whether a given MethodInstance would be compiled in the same
way now (aka, in the user's running session) as when the package was
precompiled; in practice, the way we do that is to test whether the
dispatches would be to the same methods in the current
world-age. `invoke` presents special challenges because it allows the
coder to deliberately select a different method than the one that
would be chosen by ordinary dispatch; if there is no record of how
this choice was made, it can look like it resolves to the wrong method
and this can trigger invalidation.

This allows a MethodInstance to store dispatch tuples as well as other
MethodInstances among their backedges.

Additionally:

- provide backedge-iterators for both C and Julia that abstracts
  the specific storage mechanism.

- fix a bug in the CodeInstance `relocatability` field, where methods
  that only return a constant (and hence store `nothing` for
  `inferred`) were deemed non-relocatable.

- fix a bug in which #43990 should have checked that the method had
  not been deleted. Tests passed formerly simply because we weren't
  caching external CodeInstances that inferred down to a `Const`;
  fixing that exposed the bug.  This bug has been exposed since
  merging #43990 for non-`Const` inference, and would affect Revise
  etc.

Co-authored-by: Jameson Nash <vtjnash@gmail.com>

Author: timholy

NEWS.md

@@ -35,6 +35,8 @@ Compiler/Runtime improvements
   `@nospecialize`-d call sites and avoiding excessive compilation. ([#44512])
 * All the previous usages of `@pure`-macro in `Base` has been replaced with the preferred
   `Base.@assume_effects`-based annotations. ([#44776])
+* `invoke(f, invokesig, args...)` calls to a less-specific method than would normally be chosen
+  for `f(args...)` are no longer spuriously invalidated when loading package precompile files. ([#46010])
 
 Command-line option changes
 ---------------------------

base/compiler/abstractinterpretation.jl

@@ -801,6 +801,11 @@ function collect_const_args(argtypes::Vector{Any})
                 end for i = 2:length(argtypes) ]
 end
 
+function invoke_signature(invokesig::Vector{Any})
+    ft, argtyps = widenconst(invokesig[2]), instanceof_tfunc(widenconst(invokesig[3]))[1]
+    return rewrap_unionall(Tuple{ft, unwrap_unionall(argtyps).parameters...}, argtyps)
+end
+
 function concrete_eval_call(interp::AbstractInterpreter,
     @nospecialize(f), result::MethodCallResult, arginfo::ArgInfo, sv::InferenceState)
     concrete_eval_eligible(interp, f, result, arginfo, sv) || return nothing
@@ -1631,7 +1636,7 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn
     ti = tienv[1]; env = tienv[2]::SimpleVector
     result = abstract_call_method(interp, method, ti, env, false, sv)
     (; rt, edge, effects) = result
-    edge !== nothing && add_backedge!(edge::MethodInstance, sv)
+    edge !== nothing && add_backedge!(edge::MethodInstance, sv, types)
     match = MethodMatch(ti, env, method, argtype <: method.sig)
     res = nothing
     sig = match.spec_types

base/compiler/inferencestate.jl

@@ -479,12 +479,15 @@ function add_cycle_backedge!(frame::InferenceState, caller::InferenceState, curr
 end
 
 # temporarily accumulate our edges to later add as backedges in the callee
-function add_backedge!(li::MethodInstance, caller::InferenceState)
+function add_backedge!(li::MethodInstance, caller::InferenceState, invokesig::Union{Nothing,Type}=nothing)
     isa(caller.linfo.def, Method) || return # don't add backedges to toplevel exprs
     edges = caller.stmt_edges[caller.currpc]
     if edges === nothing
         edges = caller.stmt_edges[caller.currpc] = []
     end
+    if invokesig !== nothing
+        push!(edges, invokesig)
+    end
     push!(edges, li)
     return nothing
 end

base/compiler/optimize.jl

@@ -62,6 +62,10 @@ intersect!(et::EdgeTracker, range::WorldRange) =
     et.valid_worlds[] = intersect(et.valid_worlds[], range)
 
 push!(et::EdgeTracker, mi::MethodInstance) = push!(et.edges, mi)
+function add_edge!(et::EdgeTracker, @nospecialize(invokesig), mi::MethodInstance)
+    invokesig === nothing && return push!(et.edges, mi)
+    push!(et.edges, invokesig, mi)
+end
 function push!(et::EdgeTracker, ci::CodeInstance)
     intersect!(et, WorldRange(min_world(li), max_world(li)))
     push!(et, ci.def)

base/compiler/ssair/inlining.jl

@@ -29,19 +29,21 @@ pass to apply its own inlining policy decisions.
 struct DelayedInliningSpec
     match::Union{MethodMatch, InferenceResult}
     argtypes::Vector{Any}
+    invokesig    # either nothing or a signature (signature is for an `invoke` call)
 end
+DelayedInliningSpec(match, argtypes) = DelayedInliningSpec(match, argtypes, nothing)
 
 struct InliningTodo
     # The MethodInstance to be inlined
     mi::MethodInstance
     spec::Union{ResolvedInliningSpec, DelayedInliningSpec}
 end
 
-InliningTodo(mi::MethodInstance, match::MethodMatch, argtypes::Vector{Any}) =
-    InliningTodo(mi, DelayedInliningSpec(match, argtypes))
+InliningTodo(mi::MethodInstance, match::MethodMatch, argtypes::Vector{Any}, invokesig=nothing) =
+    InliningTodo(mi, DelayedInliningSpec(match, argtypes, invokesig))
 
-InliningTodo(result::InferenceResult, argtypes::Vector{Any}) =
-    InliningTodo(result.linfo, DelayedInliningSpec(result, argtypes))
+InliningTodo(result::InferenceResult, argtypes::Vector{Any}, invokesig=nothing) =
+    InliningTodo(result.linfo, DelayedInliningSpec(result, argtypes, invokesig))
 
 struct ConstantCase
     val::Any
@@ -810,15 +812,15 @@ end
 
 function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
     mi = todo.mi
-    (; match, argtypes) = todo.spec::DelayedInliningSpec
+    (; match, argtypes, invokesig) = todo.spec::DelayedInliningSpec
     et = state.et
 
     #XXX: update_valid_age!(min_valid[1], max_valid[1], sv)
     if isa(match, InferenceResult)
         inferred_src = match.src
         if isa(inferred_src, ConstAPI)
             # use constant calling convention
-            et !== nothing && push!(et, mi)
+            et !== nothing && add_edge!(et, invokesig, mi)
             return ConstantCase(quoted(inferred_src.val))
         else
             src = inferred_src # ::Union{Nothing,CodeInfo} for NativeInterpreter
@@ -829,7 +831,7 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
         if code isa CodeInstance
             if use_const_api(code)
                 # in this case function can be inlined to a constant
-                et !== nothing && push!(et, mi)
+                et !== nothing && add_edge!(et, invokesig, mi)
                 return ConstantCase(quoted(code.rettype_const))
             else
                 src = @atomic :monotonic code.inferred
@@ -851,7 +853,7 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
 
     src === nothing && return compileable_specialization(et, match, effects)
 
-    et !== nothing && push!(et, mi)
+    et !== nothing && add_edge!(et, invokesig, mi)
     return InliningTodo(mi, retrieve_ir_for_inlining(mi, src), effects)
 end
 
@@ -873,7 +875,7 @@ function validate_sparams(sparams::SimpleVector)
     return true
 end
 
-function analyze_method!(match::MethodMatch, argtypes::Vector{Any},
+function analyze_method!(match::MethodMatch, argtypes::Vector{Any}, invokesig,
                          flag::UInt8, state::InliningState)
     method = match.method
     spec_types = match.spec_types
@@ -905,7 +907,7 @@ function analyze_method!(match::MethodMatch, argtypes::Vector{Any},
     mi = specialize_method(match; preexisting=true) # Union{Nothing, MethodInstance}
     isa(mi, MethodInstance) || return compileable_specialization(et, match, Effects())
 
-    todo = InliningTodo(mi, match, argtypes)
+    todo = InliningTodo(mi, match, argtypes, invokesig)
     # If we don't have caches here, delay resolving this MethodInstance
     # until the batch inlining step (or an external post-processing pass)
     state.mi_cache === nothing && return todo
@@ -1100,17 +1102,18 @@ function inline_invoke!(
     if isa(result, ConcreteResult)
         item = concrete_result_item(result, state)
     else
+        invokesig = invoke_signature(sig.argtypes)
         argtypes = invoke_rewrite(sig.argtypes)
         if isa(result, ConstPropResult)
-            (; mi) = item = InliningTodo(result.result, argtypes)
+            (; mi) = item = InliningTodo(result.result, argtypes, invokesig)
             validate_sparams(mi.sparam_vals) || return nothing
             if argtypes_to_type(argtypes) <: mi.def.sig
                 state.mi_cache !== nothing && (item = resolve_todo(item, state, flag))
                 handle_single_case!(ir, idx, stmt, item, todo, state.params, true)
                 return nothing
             end
         end
-        item = analyze_method!(match, argtypes, flag, state)
+        item = analyze_method!(match, argtypes, invokesig, flag, state)
     end
     handle_single_case!(ir, idx, stmt, item, todo, state.params, true)
     return nothing
@@ -1328,7 +1331,7 @@ function handle_match!(
     # during abstract interpretation: for the purpose of inlining, we can just skip
     # processing this dispatch candidate
     _any(case->case.sig === spec_types, cases) && return true
-    item = analyze_method!(match, argtypes, flag, state)
+    item = analyze_method!(match, argtypes, nothing, flag, state)
     item === nothing && return false
     push!(cases, InliningCase(spec_types, item))
     return true
@@ -1475,7 +1478,7 @@ function assemble_inline_todo!(ir::IRCode, state::InliningState)
                 if isa(result, ConcreteResult)
                     item = concrete_result_item(result, state)
                 else
-                    item = analyze_method!(info.match, sig.argtypes, flag, state)
+                    item = analyze_method!(info.match, sig.argtypes, nothing, flag, state)
                 end
                 handle_single_case!(ir, idx, stmt, item, todo, state.params)
             end

base/compiler/typeinfer.jl

@@ -312,7 +312,9 @@ function CodeInstance(
             const_flags = 0x00
         end
     end
-    relocatability = isa(inferred_result, Vector{UInt8}) ? inferred_result[end] : UInt8(0)
+    relocatability = isa(inferred_result, Vector{UInt8}) ? inferred_result[end] :
+                     inferred_result === nothing ? UInt8(1) : UInt8(0)
+    # relocatability = isa(inferred_result, Vector{UInt8}) ? inferred_result[end] : UInt8(0)
     return CodeInstance(result.linfo,
         widenconst(result_type), rettype_const, inferred_result,
         const_flags, first(valid_worlds), last(valid_worlds),
@@ -561,17 +563,12 @@ function store_backedges(frame::InferenceResult, edges::Vector{Any})
 end
 
 function store_backedges(caller::MethodInstance, edges::Vector{Any})
-    i = 1
-    while i <= length(edges)
-        to = edges[i]
+    for (typ, to) in BackedgeIterator(edges)
         if isa(to, MethodInstance)
-            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any), to, caller)
-            i += 1
+            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), to, typ, caller)
         else
             typeassert(to, Core.MethodTable)
-            typ = edges[i + 1]
             ccall(:jl_method_table_add_backedge, Cvoid, (Any, Any, Any), to, typ, caller)
-            i += 2
         end
     end
 end

base/compiler/utilities.jl

@@ -223,6 +223,59 @@ Check if `method` is declared as `Base.@constprop :none`.
 """
 is_no_constprop(method::Union{Method,CodeInfo}) = method.constprop == 0x02
 
+#############
+# backedges #
+#############
+
+"""
+    BackedgeIterator(backedges::Vector{Any})
+
+Return an iterator over a list of backedges. Iteration returns `(sig, caller)` elements,
+which will be one of the following:
+
+- `(nothing, caller::MethodInstance)`: a call made by ordinary inferrable dispatch
+- `(invokesig, caller::MethodInstance)`: a call made by `invoke(f, invokesig, args...)`
+- `(specsig, mt::MethodTable)`: an abstract call
+
+# Examples
+
+```julia
+julia> callme(x) = x+1
+callme (generic function with 1 method)
+
+julia> callyou(x) = callme(x)
+callyou (generic function with 1 method)
+
+julia> callyou(2.0)
+3.0
+
+julia> mi = first(which(callme, (Any,)).specializations)
+MethodInstance for callme(::Float64)
+
+julia> @eval Core.Compiler for (sig, caller) in BackedgeIterator(Main.mi.backedges)
+           println(sig)
+           println(caller)
+       end
+nothing
+callyou(Float64) from callyou(Any)
+```
+"""
+struct BackedgeIterator
+    backedges::Vector{Any}
+end
+
+const empty_backedge_iter = BackedgeIterator(Any[])
+
+
+function iterate(iter::BackedgeIterator, i::Int=1)
+    backedges = iter.backedges
+    i > length(backedges) && return nothing
+    item = backedges[i]
+    isa(item, MethodInstance) && return (nothing, item), i+1           # regular dispatch
+    isa(item, Core.MethodTable) && return (backedges[i+1], item), i+2  # abstract dispatch
+    return (item, backedges[i+1]::MethodInstance), i+2                 # `invoke` calls
+end
+
 #########
 # types #
 #########

base/loading.jl

@@ -2262,12 +2262,12 @@ macro __DIR__()
 end
 
 """
-    precompile(f, args::Tuple{Vararg{Any}})
+    precompile(f, argtypes::Tuple{Vararg{Any}})
 
-Compile the given function `f` for the argument tuple (of types) `args`, but do not execute it.
+Compile the given function `f` for the argument tuple (of types) `argtypes`, but do not execute it.
 """
-function precompile(@nospecialize(f), @nospecialize(args::Tuple))
-    precompile(Tuple{Core.Typeof(f), args...})
+function precompile(@nospecialize(f), @nospecialize(argtypes::Tuple))
+    precompile(Tuple{Core.Typeof(f), argtypes...})
 end
 
 const ENABLE_PRECOMPILE_WARNINGS = Ref(false)
@@ -2279,6 +2279,27 @@ function precompile(@nospecialize(argt::Type))
     return ret
 end
 
+# Variants that work for `invoke`d calls for which the signature may not be sufficient
+precompile(mi::Core.MethodInstance, world::UInt=get_world_counter()) =
+    (ccall(:jl_compile_method_instance, Cvoid, (Any, Any, UInt), mi, C_NULL, world); return true)
+
+"""
+    precompile(f, argtypes::Tuple{Vararg{Any}}, m::Method)
+
+Precompile a specific method for the given argument types. This may be used to precompile
+a different method than the one that would ordinarily be chosen by dispatch, thus
+mimicking `invoke`.
+"""
+function precompile(@nospecialize(f), @nospecialize(argtypes::Tuple), m::Method)
+    precompile(Tuple{Core.Typeof(f), argtypes...}, m)
+end
+
+function precompile(@nospecialize(argt::Type), m::Method)
+    atype, sparams = ccall(:jl_type_intersection_with_env, Any, (Any, Any), argt, m.sig)::SimpleVector
+    mi = Core.Compiler.specialize_method(m, atype, sparams)
+    return precompile(mi)
+end
+
 precompile(include_package_for_output, (PkgId, String, Vector{String}, Vector{String}, Vector{String}, typeof(_concrete_dependencies), Nothing))
 precompile(include_package_for_output, (PkgId, String, Vector{String}, Vector{String}, Vector{String}, typeof(_concrete_dependencies), String))
 precompile(create_expr_cache, (PkgId, String, String, typeof(_concrete_dependencies), IO, IO))