mock out early generator expansion in inference [ci skip]

Author: jrevels

base/inference.jl

@@ -305,7 +305,12 @@ function InferenceState(linfo::MethodInstance,
 end
 
 function get_staged(li::MethodInstance)
-    return ccall(:jl_code_for_staged, Any, (Any,), li)::CodeInfo
+    try
+        # user code might throw errors – ignore them
+        return ccall(:jl_code_for_staged, Any, (Any,), li)::CodeInfo
+    catch
+        return nothing
+    end
 end
 
 
@@ -405,12 +410,7 @@ end
 function retrieve_code_info(linfo::MethodInstance)
     m = linfo.def::Method
     if isdefined(m, :generator)
-        try
-            # user code might throw errors – ignore them
-            c = get_staged(linfo)
-        catch
-            return nothing
-        end
+        return get_staged(linfo)
     else
         # TODO: post-inference see if we can swap back to the original arrays?
         if isa(m.source, Array{UInt8,1})
@@ -1922,6 +1922,17 @@ function abstract_call_method(method::Method, @nospecialize(f), @nospecialize(si
     # Returns the topmost occurrence of that repeated edge.
     cyclei = 0
     infstate = sv
+    if isdefined(method, :generator) && method.generator.expand_early
+        method_instance = code_for_method(method, sig, sparams, sv.params.world, false)
+        if isa(method_instance, MethodInstance)
+            cinfo = get_staged(method_instance)
+            if isa(cinfo, CodeInfo) && isa(cinfo.method_for_inference_heuristics, Method)
+                checked_method = cinfo.method_for_inference_heuristics
+            end
+        end
+    else
+        checked_method = method
+    end
     while !(infstate === nothing)
         infstate = infstate::InferenceState
         if method === infstate.linfo.def && infstate.linfo.specTypes == sig
@@ -1931,7 +1942,7 @@ function abstract_call_method(method::Method, @nospecialize(f), @nospecialize(si
             break
         end
         working_method = method_for_inference_heuristics(infstate)
-        if method === working_method
+        if checked_method === working_method
             if topmost === nothing
                 # inspect the parent of this edge,
                 # to see if they are the same Method as sv