Add hook in inference recursion resolution for external AbstractInterpreter (#36401)

This extends hookability to the same-frame comparison in inference's
recursion cycle detection. The case I ran into that made this necessary
is a recursive, nested AD transform. In this case, inference must detect
if two frames have different orders of derivatives (e.g. the primitive
for `-`, again calls `-`; the external interpreter makes sure that
inference results for these end up in different caches).

Author: Keno

base/compiler/inferencestate.jl

@@ -46,6 +46,10 @@ mutable struct InferenceState
     # `max_valid`, to be used in inlining
     matching_methods_cache::IdDict{Any, Tuple{Any, UInt, UInt}}
 
+    # The interpreter that created this inference state. Not looked at by
+    # NativeInterpreter. But other interpreters may use this to detect cycles
+    interp::AbstractInterpreter
+
     # src is assumed to be a newly-allocated CodeInfo, that can be modified in-place to contain intermediate results
     function InferenceState(result::InferenceResult, src::CodeInfo,
                             cached::Bool, interp::AbstractInterpreter)
@@ -107,7 +111,8 @@ mutable struct InferenceState
             Vector{InferenceState}(), # callers_in_cycle
             #=parent=#nothing,
             cached, false, false, false,
-            IdDict{Any, Tuple{Any, UInt, UInt}}())
+            IdDict{Any, Tuple{Any, UInt, UInt}}(),
+            interp)
         result.result = frame
         cached && push!(get_inference_cache(interp), result)
         return frame

base/compiler/typeinfer.jl

@@ -439,21 +439,25 @@ function merge_call_chain!(parent::InferenceState, ancestor::InferenceState, chi
     end
 end
 
+function is_same_frame(interp::AbstractInterpreter, linfo::MethodInstance, frame::InferenceState)
+    return linfo === frame.linfo
+end
+
 # Walk through `linfo`'s upstream call chain, starting at `parent`. If a parent
 # frame matching `linfo` is encountered, then there is a cycle in the call graph
 # (i.e. `linfo` is a descendant callee of itself). Upon encountering this cycle,
 # we "resolve" it by merging the call chain, which entails unioning each intermediary
 # frame's `callers_in_cycle` field and adding the appropriate backedges. Finally,
 # we return `linfo`'s pre-existing frame. If no cycles are found, `nothing` is
 # returned instead.
-function resolve_call_cycle!(linfo::MethodInstance, parent::InferenceState)
+function resolve_call_cycle!(interp::AbstractInterpreter, linfo::MethodInstance, parent::InferenceState)
     frame = parent
     uncached = false
     limited = false
     while isa(frame, InferenceState)
         uncached |= !frame.cached # ensure we never add an uncached frame to a cycle
         limited |= frame.limited
-        if frame.linfo === linfo
+        if is_same_frame(interp, linfo, frame)
             if uncached
                 # our attempt to speculate into a constant call lead to an undesired self-cycle
                 # that cannot be converged: poison our call-stack (up to the discovered duplicate frame)
@@ -465,7 +469,7 @@ function resolve_call_cycle!(linfo::MethodInstance, parent::InferenceState)
             return frame
         end
         for caller in frame.callers_in_cycle
-            if caller.linfo === linfo
+            if is_same_frame(interp, linfo, caller)
                 if uncached
                     poison_callstack(parent, frame, false)
                     return true
@@ -496,7 +500,7 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
         # (if we asked resolve_call_cyle, it might instead detect that there is a cycle that it can't merge)
         frame = false
     else
-        frame = resolve_call_cycle!(mi, caller)
+        frame = resolve_call_cycle!(interp, mi, caller)
     end
     if frame === false
         # completely new