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

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)