apply_type_tfunc: add heuristic complexity limit (#49167)

In #48421 we removed all limits from apply_type, but that can lead to
significant expression complexity which may be hard for subtyping to
deal with. Try adding a heuristic size limiting constraint.
Also do a bunch of code rearrangement and slight cleanup so that this
change does not make the logic look complicated here.

Solves #49127

base/compiler/tfuncs.jl

@@ -1779,31 +1779,55 @@ const _tvarnames = Symbol[:_A, :_B, :_C, :_D, :_E, :_F, :_G, :_H, :_I, :_J, :_K,
  uncertain = true
  unw = unwrap_unionall(ai)
  isT = isType(unw)
+ # compute our desired upper bound value
  if isT
- tai = ai
- while isa(tai, UnionAll)
- if contains_is(outervars, tai.var)
- ai = rename_unionall(ai)
- unw = unwrap_unionall(ai)
- break
+ ub = rewrap_unionall(unw.parameters[1], ai)
+ else
+ ub = Any
+ end
+ if !istuple && unionall_depth(ai) > 3
+ # Heuristic: if we are adding more than N unknown parameters here to the
+ # outer type, use the wrapper type, instead of letting it nest more
+ # complexity here. This is not monotonic, but seems to work out pretty well.
+ if isT
+ ub = unwrap_unionall(unw.parameters[1])
+ if ub isa DataType
+ ub = ub.name.wrapper
+ unw = Type{unwrap_unionall(ub)}
+ ai = rewrap_unionall(unw, ub)
+ else
+ isT = false
+ ai = unw = ub = Any
  end
- tai = tai.body
+ else
+ isT = false
+ ai = unw = ub = Any
  end
+ elseif !isT
+ # if we didn't have isType to compute ub directly, try to use instanceof_tfunc to refine this guess
+ ai_w = widenconst(ai)
+ ub = ai_w isa Type && ai_w <: Type ? instanceof_tfunc(ai)[1] : Any
  end
- ai_w = widenconst(ai)
- ub = ai_w isa Type && ai_w <: Type ? instanceof_tfunc(ai)[1] : Any
  if istuple
  # in the last parameter of a Tuple type, if the upper bound is Any
  # then this could be a Vararg type.
  if i == largs && ub === Any
- push!(tparams, Vararg)
- elseif isT
- push!(tparams, rewrap_unionall((unw::DataType).parameters[1], ai))
- else
- push!(tparams, Any)
+ ub = Vararg
  end
+ push!(tparams, ub)
  elseif isT
- push!(tparams, (unw::DataType).parameters[1])
+ tai = ai
+ while isa(tai, UnionAll)
+ # make sure vars introduced here are unique
+ if contains_is(outervars, tai.var)
+ ai = rename_unionall(ai)
+ unw = unwrap_unionall(ai)::DataType
+ # ub = rewrap_unionall(unw, ai)
+ break
+ end
+ tai = tai.body
+ end
+ push!(tparams, unw.parameters[1])
  while isa(ai, UnionAll)
  push!(outervars, ai.var)
  ai = ai.body

base/compiler/typeutils.jl

@@ -308,6 +308,15 @@ function _unioncomplexity(@nospecialize x)
  end
 end
 
+function unionall_depth(@nospecialize ua) # aka subtype_env_size
+ depth = 0
+ while ua isa UnionAll
+ depth += 1
+ ua = ua.body
+ end
+ return depth
+end
+
 # convert a Union of Tuple types to a Tuple of Unions
 function unswitchtupleunion(u::Union)
  ts = uniontypes(u)

test/compiler/inference.jl

@@ -2658,10 +2658,16 @@ let 𝕃 = Core.Compiler.fallback_lattice
  @test apply_type_tfunc(𝕃, Const(Issue47089), Const(String)) === Union{}
  @test apply_type_tfunc(𝕃, Const(Issue47089), Const(AbstractString)) === Union{}
  @test apply_type_tfunc(𝕃, Const(Issue47089), Type{Ptr}, Type{Ptr{T}} where T) === Base.rewrap_unionall(Type{Issue47089.body.body}, Issue47089)
+ # check complexity size limiting
+ @test apply_type_tfunc(𝕃, Const(Val), Type{Pair{Pair{Pair{Pair{A,B},C},D},E}} where {A,B,C,D,E}) == Type{Val{Pair{A, B}}} where {A, B}
+ @test apply_type_tfunc(𝕃, Const(Pair), Base.rewrap_unionall(Type{Pair.body.body},Pair), Type{Pair{Pair{Pair{Pair{A,B},C},D},E}} where {A,B,C,D,E}) == Type{Pair{Pair{A, B}, Pair{C, D}}} where {A, B, C, D}
+ @test apply_type_tfunc(𝕃, Const(Val), Type{Union{Int,Pair{Pair{Pair{Pair{A,B},C},D},E}}} where {A,B,C,D,E}) == Type{Val{_A}} where _A
 end
 @test only(Base.return_types(keys, (Dict{String},))) == Base.KeySet{String, T} where T<:(Dict{String})
 @test only(Base.return_types((r)->similar(Array{typeof(r[])}, 1), (Base.RefValue{Array{Int}},))) == Vector{<:Array{Int}}
 @test only(Base.return_types((r)->similar(Array{typeof(r[])}, 1), (Base.RefValue{Array{<:Real}},))) == Vector{<:Array{<:Real}}
+# test complexity limit on apply_type on a function capturing functions returning functions
+@test only(Base.return_types(Base.afoldl, (typeof((m, n) -> () -> Returns(nothing)(m, n)), Function, Function, Vararg{Function}))) === Function
 
 let A = Tuple{A,B,C,D,E,F,G,H} where {A,B,C,D,E,F,G,H}
  B = Core.Compiler.rename_unionall(A)