fix the optimizer's use of cached varargs type info for _apply unrolling

Author: jrevels

base/compiler/ssair/inlining.jl

@@ -557,22 +557,43 @@ function spec_lambda(@nospecialize(atype), sv::OptimizationState, @nospecialize(
     linfo
 end
 
-function rewrite_apply_exprargs!(ir::IRCode, idx::Int, argexprs::Vector{Any}, sv::OptimizationState)
+function rewrite_apply_exprargs!(ir::IRCode, idx::Int, argexprs::Vector{Any}, atypes::Vector{Any}, sv::OptimizationState)
     new_argexprs = Any[argexprs[2]]
-    # Flatten all tuples
-    for arg in argexprs[3:end]
-        tupT = argextype(arg, ir, sv.sp)
-        t = widenconst(tupT)
-        for i = 1:length(t.parameters)
-            # Insert a getfield call here
-            new_call = Expr(:call, Core.getfield, arg, i)
-            typ = getfield_tfunc(tupT, Const(i))
-            new_arg = insert_node!(ir, idx, typ, new_call)
-            push!(new_argexprs, new_arg)
+    new_atypes = Any[atypes[2]]
+    # loop over original arguments and flatten any known iterators
+    for i in 3:length(argexprs)
+        def = argexprs[i]
+        # As a special case, if we can see the tuple() call, look at it's arguments to find
+        # our types. They can be more precise (e.g. f(Bool, A...) would be lowered as
+        # _apply(f, tuple(Bool)::Tuple{DataType}, A), which might not be precise enough to
+        # get a good method match). This pattern is used in the array code a bunch.
+        if isa(def, SSAValue) && is_tuple_call(ir, ir[def])
+            def_args = ir[def].args
+            def_atypes = Any[argextype(def_args[i], ir, sv.sp) for i in 2:length(def_args)]
+        elseif isa(def, Argument) && def.n === length(ir.argtypes) && !isempty(sv.result_vargs)
+            def_atypes = sv.result_vargs
+        else
+            def_atypes = Any[]
+            for p in widenconst(atypes[i]).parameters
+                if isa(p, DataType) && isdefined(p, :instance)
+                    # replace singleton types with their equivalent Const object
+                    p = Const(p.instance)
+                elseif isconstType(p)
+                    p = Const(p.parameters[1])
+                end
+                push!(def_atypes, p)
+            end
+        end
+        # now push flattened types into new_atypes and getfield exprs into new_argexprs
+        for j in 1:length(def_atypes)
+            def_atype = def_atypes[j]
+            new_call = Expr(:call, Core.getfield, def, j)
+            new_argexpr = insert_node!(ir, idx, def_atype, new_call)
+            push!(new_argexprs, new_argexpr)
+            push!(new_atypes, def_atype)
         end
     end
-    argexprs = new_argexprs
-    return argexprs
+    return new_argexprs, new_atypes
 end
 
 function rewrite_invoke_exprargs!(inserter, argexprs::Vector{Any})
@@ -780,52 +801,27 @@ function assemble_inline_todo!(ir::IRCode, linetable::Vector{LineInfoNode}, sv::
         isapply = isinvoke = false
 
         # Handle _apply
-        isapply = false
         if f === Core._apply
-            new_atypes = Any[]
-            ft = argextype(stmt.args[2], ir, sv.sp)
+            ft = atypes[2]
             has_free_typevars(ft) && continue
             f = singleton_type(ft)
-            # Push function type
-            push!(new_atypes, ft)
             # Try to figure out the signature of the function being called
+            # and if rewrite_apply_exprargs can deal with this form
             ok = true
-            for (typ, def) in zip(atypes[3:end], stmt.args[3:end])
-                typ = widenconst(typ)
-                # We don't know what'll be in a SimpleVector, bail out
-                isa(typ, SimpleVector) && (ok = false; break)
+            for i = 3:length(atypes)
+                typ = widenconst(atypes[i])
                 # TODO: We could basically run the iteration protocol here
-                if !isa(typ, DataType) || typ.name !== Tuple.name || isvatuple(typ)
+                if !isa(typ, DataType) || typ.name !== Tuple.name ||
+                    isvatuple(typ) || length(typ.parameters) > sv.params.MAX_TUPLE_SPLAT
                     ok = false
                     break
                 end
-                if length(typ.parameters) > sv.params.MAX_TUPLE_SPLAT
-                    ok = false
-                    break
-                end
-                # As a special case, if we can see the tuple() call, look at it's arguments to find
-                # our types. They can be more precise (e.g. f(Bool, A...) would be lowered as
-                # _apply(f, tuple(Bool)::Tuple{DataType}, A), which might not be precise enough to
-                # get a good method match). This pattern is used in the array code a bunch.
-                if isa(def, SSAValue) && is_tuple_call(ir, ir[def])
-                    for tuparg in ir[def].args[2:end]
-                        push!(new_atypes, argextype(tuparg, ir, sv.sp))
-                    end
-                elseif isa(def, Argument) && def.n === length(ir.argtypes) && !isempty(sv.result_vargs)
-                    append!(new_atypes, sv.result_vargs)
-                else
-                    append!(new_atypes, typ.parameters)
-                end
             end
             ok || continue
-            atypes = new_atypes
             isapply = true
-        end
-
-        # Independent of whether we can inline, the above analysis allows us to rewrite
-        # this apply call to a regular call
-        if isapply
-            stmt.args = rewrite_apply_exprargs!(ir, idx, stmt.args, sv)
+            # Independent of whether we can inline, the above analysis allows us to rewrite
+            # this apply call to a regular call
+            stmt.args, atypes = rewrite_apply_exprargs!(ir, idx, stmt.args, atypes, sv)
         end
 
         if f !== Core.invoke && (isa(f, IntrinsicFunction) || ft ⊑ IntrinsicFunction || isa(f, Builtin) || ft ⊑ Builtin)