AbsInt: add interfaces to customize cases when cached results are used (

#53318)

For external abstract interpreters like JET, which propagate custom data
interprocedurally, it is essential to inject custom behaviors into where
cached regular/constant inference results are used. Previously, this was
accomplished by overloading both `abstract_call_method` and
`get(::WorldView{CustomView}, ...)` (and `const_prop_call` and
`cached_lookup`), that method was admittedly hacky and should probably
better to be avoided. Moreover, after #52233, doing so has become
infeasible when the external abstract interpreter uses
`InternalCodeCache`.

To address this issue, this commit provides an interface named
`return_cached_result`. This allows external abstract interpreters to
inject custom interprocedural data propagation during abstract
interpretation even when they use `InternalCodeCache`.

base/compiler/abstractinterpretation.jl

@@ -1221,45 +1221,53 @@ function semi_concrete_eval_call(interp::AbstractInterpreter,
     return nothing
 end
 
+const_prop_result(inf_result::InferenceResult) =
+    ConstCallResults(inf_result.result, inf_result.exc_result, ConstPropResult(inf_result),
+                     inf_result.ipo_effects, inf_result.linfo)
+
+# return cached constant analysis result
+return_cached_result(::AbstractInterpreter, inf_result::InferenceResult, ::AbsIntState) =
+    const_prop_result(inf_result)
+
 function const_prop_call(interp::AbstractInterpreter,
     mi::MethodInstance, result::MethodCallResult, arginfo::ArgInfo, sv::AbsIntState,
     concrete_eval_result::Union{Nothing, ConstCallResults}=nothing)
     inf_cache = get_inference_cache(interp)
     𝕃ᵢ = typeinf_lattice(interp)
     inf_result = cache_lookup(𝕃ᵢ, mi, arginfo.argtypes, inf_cache)
-    if inf_result === nothing
-        # fresh constant prop'
-        argtypes = has_conditional(𝕃ᵢ, sv) ? ConditionalArgtypes(arginfo, sv) : SimpleArgtypes(arginfo.argtypes)
-        inf_result = InferenceResult(mi, argtypes, typeinf_lattice(interp))
-        if !any(inf_result.overridden_by_const)
-            add_remark!(interp, sv, "[constprop] Could not handle constant info in matching_cache_argtypes")
-            return nothing
-        end
-        frame = InferenceState(inf_result, #=cache_mode=#:local, interp)
-        if frame === nothing
-            add_remark!(interp, sv, "[constprop] Could not retrieve the source")
-            return nothing # this is probably a bad generated function (unsound), but just ignore it
-        end
-        frame.parent = sv
-        if !typeinf(interp, frame)
-            add_remark!(interp, sv, "[constprop] Fresh constant inference hit a cycle")
-            return nothing
-        end
-        @assert inf_result.result !== nothing
-        if concrete_eval_result !== nothing
-            # override return type and effects with concrete evaluation result if available
-            inf_result.result = concrete_eval_result.rt
-            inf_result.ipo_effects = concrete_eval_result.effects
-        end
-    else
+    if inf_result !== nothing
         # found the cache for this constant prop'
         if inf_result.result === nothing
             add_remark!(interp, sv, "[constprop] Found cached constant inference in a cycle")
             return nothing
         end
+        @assert inf_result.linfo === mi "MethodInstance for cached inference result does not match"
+        return return_cached_result(interp, inf_result, sv)
+    end
+    # perform fresh constant prop'
+    argtypes = has_conditional(𝕃ᵢ, sv) ? ConditionalArgtypes(arginfo, sv) : SimpleArgtypes(arginfo.argtypes)
+    inf_result = InferenceResult(mi, argtypes, typeinf_lattice(interp))
+    if !any(inf_result.overridden_by_const)
+        add_remark!(interp, sv, "[constprop] Could not handle constant info in matching_cache_argtypes")
+        return nothing
+    end
+    frame = InferenceState(inf_result, #=cache_mode=#:local, interp)
+    if frame === nothing
+        add_remark!(interp, sv, "[constprop] Could not retrieve the source")
+        return nothing # this is probably a bad generated function (unsound), but just ignore it
+    end
+    frame.parent = sv
+    if !typeinf(interp, frame)
+        add_remark!(interp, sv, "[constprop] Fresh constant inference hit a cycle")
+        return nothing
+    end
+    @assert inf_result.result !== nothing
+    if concrete_eval_result !== nothing
+        # override return type and effects with concrete evaluation result if available
+        inf_result.result = concrete_eval_result.rt
+        inf_result.ipo_effects = concrete_eval_result.effects
     end
-    return ConstCallResults(inf_result.result, inf_result.exc_result,
-        ConstPropResult(inf_result), inf_result.ipo_effects, mi)
+    return const_prop_result(inf_result)
 end
 
 # TODO implement MustAlias forwarding

base/compiler/typeinfer.jl

@@ -818,23 +818,29 @@ struct EdgeCallResult
     end
 end
 
+# return cached regular inference result
+function return_cached_result(::AbstractInterpreter, codeinst::CodeInstance, caller::AbsIntState)
+    rt = cached_return_type(codeinst)
+    effects = ipo_effects(codeinst)
+    update_valid_age!(caller, WorldRange(min_world(codeinst), max_world(codeinst)))
+    return EdgeCallResult(rt, codeinst.exctype, codeinst.def, effects)
+end
+
 # compute (and cache) an inferred AST and return the current best estimate of the result type
 function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize(atype), sparams::SimpleVector, caller::AbsIntState)
     mi = specialize_method(method, atype, sparams)::MethodInstance
-    code = get(code_cache(interp), mi, nothing)
+    codeinst = get(code_cache(interp), mi, nothing)
     force_inline = is_stmt_inline(get_curr_ssaflag(caller))
-    if code isa CodeInstance # return existing rettype if the code is already inferred
-        inferred = @atomic :monotonic code.inferred
+    if codeinst isa CodeInstance # return existing rettype if the code is already inferred
+        inferred = @atomic :monotonic codeinst.inferred
         if inferred === nothing && force_inline
             # we already inferred this edge before and decided to discard the inferred code,
             # nevertheless we re-infer it here again in order to propagate the re-inferred
             # source to the inliner as a volatile result
             cache_mode = CACHE_MODE_VOLATILE
         else
-            rt = cached_return_type(code)
-            effects = ipo_effects(code)
-            update_valid_age!(caller, WorldRange(min_world(code), max_world(code)))
-            return EdgeCallResult(rt, code.exctype, mi, effects)
+            @assert codeinst.def === mi "MethodInstance for cached edge does not match"
+            return return_cached_result(interp, codeinst, caller)
         end
     else
         cache_mode = CACHE_MODE_GLOBAL # cache edge targets globally by default