add method_for_inference_heuristics field to CodeInfo

Author: jrevels

base/inference.jl

@@ -1909,6 +1909,11 @@ function abstract_call_gf_by_type(@nospecialize(f), @nospecialize(atype), sv::In
     return rettype
 end
 
+function method_for_inference_heuristics(infstate::InferenceState)
+    m = infstate.src.method_for_inference_heuristics
+    return m === nothing ? infstate.linfo.def : m
+end
+
 function abstract_call_method(method::Method, @nospecialize(f), @nospecialize(sig), sparams::SimpleVector, sv::InferenceState)
     topmost = nothing
     # Limit argument type tuple growth of functions:
@@ -1919,13 +1924,14 @@ function abstract_call_method(method::Method, @nospecialize(f), @nospecialize(si
     infstate = sv
     while !(infstate === nothing)
         infstate = infstate::InferenceState
-        if method === infstate.linfo.def
-            if infstate.linfo.specTypes == sig
-                # avoid widening when detecting self-recursion
-                # TODO: merge call cycle and return right away
-                topmost = nothing
-                break
-            end
+        if method === infstate.linfo.def && infstate.linfo.specTypes == sig
+            # avoid widening when detecting self-recursion
+            # TODO: merge call cycle and return right away
+            topmost = nothing
+            break
+        end
+        working_method = method_for_inference_heuristics(infstate)
+        if method === working_method
             if topmost === nothing
                 # inspect the parent of this edge,
                 # to see if they are the same Method as sv
@@ -1943,7 +1949,8 @@ function abstract_call_method(method::Method, @nospecialize(f), @nospecialize(si
                     # then check the parent link
                     if topmost === nothing && parent !== nothing
                         parent = parent::InferenceState
-                        if parent.cached && parent.linfo.def === sv.linfo.def
+                        parent_method = method_for_inference_heuristics(parent)
+                        if parent.cached && parent_method === working_method
                             topmost = infstate
                         end
                     end
@@ -3220,6 +3227,7 @@ function typeinf_code(linfo::MethodInstance, optimize::Bool, cached::Bool,
                     tree.slotflags = UInt8[ 0 for i = 1:method.nargs ]
                     tree.slottypes = nothing
                     tree.ssavaluetypes = 0
+                    tree.method_for_inference_heuristics = nothing
                     tree.inferred = true
                     tree.pure = true
                     tree.inlineable = true

src/jltypes.c

@@ -2044,8 +2044,9 @@ void jl_init_types(void)
     jl_code_info_type =
         jl_new_datatype(jl_symbol("CodeInfo"), core,
                         jl_any_type, jl_emptysvec,
-                        jl_perm_symsvec(9,
+                        jl_perm_symsvec(10,
                             "code",
+                            "method_for_inference_heuristics",
                             "slottypes",
                             "ssavaluetypes",
                             "slotflags",
@@ -2054,17 +2055,18 @@ void jl_init_types(void)
                             "inlineable",
                             "propagate_inbounds",
                             "pure"),
-                        jl_svec(9,
+                        jl_svec(10,
                             jl_array_any_type,
                             jl_any_type,
                             jl_any_type,
+                            jl_any_type,
                             jl_array_uint8_type,
                             jl_array_any_type,
                             jl_bool_type,
                             jl_bool_type,
                             jl_bool_type,
                             jl_bool_type),
-                        0, 1, 9);
+                        0, 1, 10);
 
     jl_method_type =
         jl_new_datatype(jl_symbol("Method"), core,

src/julia.h

@@ -229,6 +229,7 @@ typedef struct _jl_llvm_functions_t {
 // This type describes a single function body
 typedef struct _jl_code_info_t {
     jl_array_t *code;  // Any array of statements
+    jl_value_t *method_for_inference_heuristics; // optional method used during inference
     jl_value_t *slottypes; // types of variable slots (or `nothing`)
     jl_value_t *ssavaluetypes;  // types of ssa values (or count of them)
     jl_array_t *slotflags;  // local var bit flags

src/method.c

@@ -157,6 +157,7 @@ static void jl_code_info_set_ast(jl_code_info_t *li, jl_expr_t *ast)
                 jl_array_del_end(meta, na - ins);
         }
     }
+    li->method_for_inference_heuristics = jl_nothing;
     jl_array_t *vinfo = (jl_array_t*)jl_exprarg(ast, 1);
     jl_array_t *vis = (jl_array_t*)jl_array_ptr_ref(vinfo, 0);
     size_t nslots = jl_array_len(vis);
@@ -225,6 +226,7 @@ JL_DLLEXPORT jl_code_info_t *jl_new_code_info_uninit(void)
         (jl_code_info_t*)jl_gc_alloc(ptls, sizeof(jl_code_info_t),
                                        jl_code_info_type);
     src->code = NULL;
+    src->method_for_inference_heuristics = NULL;
     src->slotnames = NULL;
     src->slotflags = NULL;
     src->slottypes = NULL;

src/toplevel.c

@@ -548,6 +548,7 @@ static jl_code_info_t *expr_to_code_info(jl_value_t *expr)
     jl_gc_wb(src, src->slotflags);
     src->ssavaluetypes = jl_box_long(0);
     jl_gc_wb(src, src->ssavaluetypes);
+    src->method_for_inference_heuristics = jl_nothing;
 
     JL_GC_POP();
     return src;