use type signature instead of Method for inference-spoofing CodeInfo field [ci skip]

Author: jrevels

base/inference.jl

@@ -2085,19 +2085,28 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
     return result
 end
 
-function method_for_inference_heuristics(infstate::InferenceState)
-    m = infstate.src.method_for_inference_heuristics
-    return isa(m, Method) ? m : infstate.linfo.def
+function method_for_inference_heuristics(cinfo, default::Method)::Method
+    if isa(cinfo, CodeInfo)
+        # appropriate format for `sig` is svec(ftype, argtypes, world)
+        sig = cinfo.signature_for_inference_heuristics
+        if isa(sig, SimpleVector) && length(sig) == 3
+            methods = _methods(sig[1], sig[2], -1, sig[3])
+            if length(methods) == 1
+                _, _, m = methods[]
+                if isa(m, Method)
+                    return m
+                end
+            end
+        end
+    end
+    return default
 end
 
-function method_for_inference_heuristics(method::Method, @nospecialize(sig), sparams, world)
+function method_for_inference_heuristics(method::Method, @nospecialize(sig), sparams, world)::Method
     if isdefined(method, :generator) && method.generator.expand_early
         method_instance = code_for_method(method, sig, sparams, world, false)
         if isa(method_instance, MethodInstance)
-            cinfo = get_staged(method_instance)
-            if isa(cinfo, CodeInfo) && isa(cinfo.method_for_inference_heuristics, Method)
-                return cinfo.method_for_inference_heuristics
-            end
+            return method_for_inference_heuristics(get_staged(method_instance), method)
         end
     end
     return method
@@ -2122,7 +2131,7 @@ function abstract_call_method(method::Method, @nospecialize(sig), sparams::Simpl
             edgecycle = true
             break
         end
-        working_method = method_for_inference_heuristics(infstate)
+        working_method = method_for_inference_heuristics(infstate.src, infstate.linfo.def)
         if checked_method === working_method
             if topmost === nothing
                 # inspect the parent of this edge,
@@ -2142,7 +2151,7 @@ function abstract_call_method(method::Method, @nospecialize(sig), sparams::Simpl
                     # then check the parent link
                     if topmost === nothing && parent !== nothing
                         parent = parent::InferenceState
-                        parent_method = method_for_inference_heuristics(parent)
+                        parent_method = method_for_inference_heuristics(parent.src, parent.linfo.def)
                         if parent.cached && parent_method === working_method
                             topmost = infstate
                             edgecycle = true
@@ -3408,7 +3417,7 @@ function typeinf_code(linfo::MethodInstance, optimize::Bool, cached::Bool,
                     method = linfo.def::Method
                     tree = ccall(:jl_new_code_info_uninit, Ref{CodeInfo}, ())
                     tree.code = Any[ Expr(:return, quoted(linfo.inferred_const)) ]
-                    tree.method_for_inference_heuristics = nothing
+                    tree.signature_for_inference_heuristics = nothing
                     tree.slotnames = Any[ compiler_temp_sym for i = 1:method.nargs ]
                     tree.slotflags = UInt8[ 0 for i = 1:method.nargs ]
                     tree.slottypes = nothing

src/jltypes.c

@@ -2046,7 +2046,7 @@ void jl_init_types(void)
                         jl_any_type, jl_emptysvec,
                         jl_perm_symsvec(10,
                             "code",
-                            "method_for_inference_heuristics"
+                            "signature_for_inference_heuristics"
                             "slottypes",
                             "ssavaluetypes",
                             "slotflags",

src/julia.h

@@ -229,7 +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 *signature_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

@@ -187,7 +187,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;
+    li->signature_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);
@@ -256,7 +256,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->signature_for_inference_heuristics = NULL;
     src->slotnames = NULL;
     src->slotflags = NULL;
     src->slottypes = NULL;

src/toplevel.c

@@ -545,7 +545,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;
+    src->signature_for_inference_heuristics = jl_nothing;
 
     JL_GC_POP();
     return src;

test/inference.jl

@@ -1324,7 +1324,7 @@ end
 
 function f24852_gen_cinfo_inflated(X, Y, f, x, y)
     method, code_info = f24852_kernel_cinfo(x, y)
-    code_info.method_for_inference_heuristics = method
+    code_info.signature_for_inference_heuristics = Core.Inference.svec(f, (x, y), typemax(UInt))
     return code_info
 end
 
@@ -1365,7 +1365,7 @@ end
 x, y = rand(), rand()
 result = f24852_kernel(x, y)
 
-# TODO: The commented out tests here are the ones where `method_for_inference_heuristics`
+# TODO: The commented out tests here are the ones where `signature_for_inference_heuristics`
 # is inflated; these tests cause segfaults. Probably due to incorrect CodeInfo
 # construction/initialization happening somewhere...
 
@@ -1377,5 +1377,5 @@ result = f24852_kernel(x, y)
 @test result === f24852_early_uninflated(x, y)
 # @test result === f24852_early_inflated(x, y)
 
-# TODO: test that `expand_early = true` + inflated `method_for_inference_heuristics`
+# TODO: test that `expand_early = true` + inflated `signature_for_inference_heuristics`
 # can be used to tighten up some inference result.