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

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 JuliaLang#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