inference: first try to tmerge tuples to a single Vararg type (#22120)

base/compiler/typelimits.jl

@@ -305,7 +305,7 @@ end
 # pick a wider type that contains both typea and typeb,
 # with some limits on how "large" it can get,
 # but without losing too much precision in common cases
-# and also trying to be associative and commutative
+# and also trying to be mostly associative and commutative
 function tmerge(@nospecialize(typea), @nospecialize(typeb))
  typea ⊑ typeb && return typeb
  typeb ⊑ typea && return typea
@@ -384,15 +384,14 @@ function tmerge(@nospecialize(typea), @nospecialize(typeb))
  types[j] = Union{}
  typenames[j] = Any.name
  else
- widen = typenames[i].wrapper
  if typenames[i] === Tuple.name
  # try to widen Tuple slower: make a single non-concrete Tuple containing both
  # converge the Tuple element-wise if they are the same length
  # see 4ee2b41552a6bc95465c12ca66146d69b354317b, be59686f7613a2ccfd63491c7b354d0b16a95c05,
- if nothing !== tuplelen(ti) === tuplelen(tj)
-  widen = tuplemerge(ti, tj)
-  end
- # TODO: else, try to merge them into a single Tuple{Vararg{T}} instead (#22120)?
+ widen = tuplemerge(unwrap_unionall(ti)::DataType, unwrap_unionall(tj)::DataType)
+ widen = rewrap_unionall(rewrap_unionall(widen, ti), tj)
+ else
+ widen = typenames[i].wrapper
  end
  types[i] = Union{}
  types[j] = widen
@@ -411,30 +410,77 @@ function tmerge(@nospecialize(typea), @nospecialize(typeb))
 end
 
 # the inverse of switchtupleunion, with limits on max element union size
-function tuplemerge(@nospecialize(a), @nospecialize(b))
- if isa(a, UnionAll)
- return UnionAll(a.var, tuplemerge(a.body, b))
- elseif isa(b, UnionAll)
- return UnionAll(b.var, tuplemerge(a, b.body))
- elseif isa(a, Union)
- return tuplemerge(tuplemerge(a.a, a.b), b)
- elseif isa(b, Union)
- return tuplemerge(a, tuplemerge(b.a, b.b))
- end
- a = a::DataType
- b = b::DataType
+function tuplemerge(a::DataType, b::DataType)
+ @assert a.name === b.name === Tuple.name "assertion failure"
  ap, bp = a.parameters, b.parameters
  lar = length(ap)::Int
  lbr = length(bp)::Int
- @assert lar === lbr && a.name === b.name === Tuple.name "assertion failure"
- p = Vector{Any}(undef, lar)
- for i = 1:lar
+ va = lar > 0 && isvarargtype(ap[lar])
+ vb = lbr > 0 && isvarargtype(bp[lbr])
+ if lar == lbr && !va && !vb
+ lt = lar
+ vt = false
+ else
+ lt = 0 # or min(lar - va, lbr - vb)
+ vt = true
+ end
+ # combine the common elements
+ p = Vector{Any}(undef, lt + vt)
+ for i = 1:lt
  ui = Union{ap[i], bp[i]}
  if unionlen(ui) < MAX_TYPEUNION_LEN
  p[i] = ui
  else
  p[i] = Any
  end
  end
+ # merge the remaining tail into a single, simple Tuple{Vararg{T}} (#22120)
+ if vt
+ tail = Union{}
+ for loop_b = (false, true)
+ for i = (lt + 1):(loop_b ? lbr : lar)
+ ti = unwrapva(loop_b ? bp[i] : ap[i])
+ # compare (ti <-> tail), (wrapper ti <-> tail), (ti <-> wrapper tail), then (wrapper ti <-> wrapper tail)
+ # until we find the first element that contains the other in the pair
+ # TODO: this result would be more stable (and more associative and more commutative)
+ # if we either joined all of the element wrappers first into a wide-tail, then picked between that or an exact tail,
+ # or (equivalently?) iteratively took super-types until reaching a common wrapper
+ # e.g. consider the results of `tuplemerge(Tuple{Complex}, Tuple{Number, Int})` and of
+ # `tuplemerge(Tuple{Int}, Tuple{String}, Tuple{Int, String})`
+ if !(ti <: tail)
+ if tail <: ti
+ tail = ti # widen to ti
+ else
+ uw = unwrap_unionall(tail)
+ if uw isa DataType && tail <: uw.name.wrapper
+ # widen tail to wrapper(tail)
+ tail = uw.name.wrapper
+ if !(ti <: tail)
+ #assert !(tail <: ti)
+ uw = unwrap_unionall(ti)
+ if uw isa DataType && ti <: uw.name.wrapper
+ # widen ti to wrapper(ti)
+ ti = uw.name.wrapper
+ #assert !(ti <: tail)
+ if tail <: ti
+ tail = ti
+ else
+ tail = Any # couldn't find common super-type
+ end
+ else
+ tail = Any # couldn't analyze type
+ end
+ end
+ else
+ tail = Any # couldn't analyze type
+ end
+ end
+ end
+ tail === Any && return Tuple # short-circuit loop
+ end
+ end
+ @assert !(tail === Union{})
+ p[lt + 1] = Vararg{tail}
+ end
  return Tuple{p...}
 end

base/compiler/typeutils.jl

@@ -139,23 +139,3 @@ function _switchtupleunion(t::Vector{Any}, i::Int, tunion::Vector{Any}, @nospeci
  end
  return tunion
 end
-
-tuplelen(@nospecialize tpl) = nothing
-function tuplelen(tpl::DataType)
- l = length(tpl.parameters)::Int
- if l > 0
- last = unwrap_unionall(tpl.parameters[l])
- if isvarargtype(last)
- N = last.parameters[2]
- N isa Int || return nothing
- l += N - 1
- end
- end
- return l
-end
-tuplelen(tpl::UnionAll) = tuplelen(tpl.body)
-function tuplelen(tpl::Union)
- la, lb = tuplelen(tpl.a), tuplelen(tpl.b)
- la == lb && return la
- return nothing
-end

base/essentials.jl

@@ -204,21 +204,23 @@ end
 const _va_typename = Vararg.body.body.name
 function isvarargtype(@nospecialize(t))
  t = unwrap_unionall(t)
- isa(t, DataType) && (t::DataType).name === _va_typename
+ return isa(t, DataType) && (t::DataType).name === _va_typename
 end
 
 isvatuple(t::DataType) = (n = length(t.parameters); n > 0 && isvarargtype(t.parameters[n]))
 function unwrapva(@nospecialize(t))
+ # NOTE: this returns a related type, but it's NOT a subtype of the original tuple
  t2 = unwrap_unionall(t)
- isvarargtype(t2) ? rewrap_unionall(t2.parameters[1],t) : t
+ return isvarargtype(t2) ? rewrap_unionall(t2.parameters[1], t) : t
 end
 
 typename(a) = error("typename does not apply to this type")
 typename(a::DataType) = a.name
 function typename(a::Union)
  ta = typename(a.a)
  tb = typename(a.b)
- ta === tb ? tb : error("typename does not apply to unions whose components have different typenames")
+ ta === tb || error("typename does not apply to unions whose components have different typenames")
+ return tb
 end
 typename(union::UnionAll) = typename(union.body)
 

test/compiler/compiler.jl

@@ -18,6 +18,48 @@ let comparison = Tuple{X, X} where X<:Tuple
  @test Core.Compiler.limit_type_size(sig, ref, Tuple{comparison}, 100, 10) == sig
 end
 
+# PR 22120
+function tmerge_test(a, b, r, commutative=true)
+ @test r == Core.Compiler.tuplemerge(a, b)
+ if commutative
+ @test r == Core.Compiler.tuplemerge(b, a)
+ else
+ @test_broken r == Core.Compiler.tuplemerge(b, a)
+ end
+end
+tmerge_test(Tuple{Int}, Tuple{String}, Tuple{Union{Int, String}})
+tmerge_test(Tuple{Int}, Tuple{String, String}, Tuple)
+tmerge_test(Tuple{Vararg{Int}}, Tuple{String}, Tuple)
+tmerge_test(Tuple{Int}, Tuple{Int, Int},
+ Tuple{Vararg{Int}})
+tmerge_test(Tuple{Integer}, Tuple{Int, Int},
+ Tuple{Vararg{Integer}})
+tmerge_test(Tuple{}, Tuple{Int, Int},
+ Tuple{Vararg{Int}})
+tmerge_test(Tuple{}, Tuple{Complex},
+ Tuple{Vararg{Complex}})
+tmerge_test(Tuple{ComplexF32}, Tuple{ComplexF32, ComplexF64},
+ Tuple{Vararg{Complex}})
+tmerge_test(Tuple{Vararg{ComplexF32}}, Tuple{Vararg{ComplexF64}},
+ Tuple{Vararg{Complex}})
+tmerge_test(Tuple{}, Tuple{ComplexF32, Vararg{Union{ComplexF32, ComplexF64}}},
+ Tuple{Vararg{Union{ComplexF32, ComplexF64}}})
+tmerge_test(Tuple{ComplexF32}, Tuple{ComplexF32, Vararg{Union{ComplexF32, ComplexF64}}},
+ Tuple{Vararg{Union{ComplexF32, ComplexF64}}})
+tmerge_test(Tuple{ComplexF32, ComplexF32, ComplexF32}, Tuple{ComplexF32, Vararg{Union{ComplexF32, ComplexF64}}},
+ Tuple{Vararg{Union{ComplexF32, ComplexF64}}})
+tmerge_test(Tuple{}, Tuple{Union{ComplexF64, ComplexF32}, Vararg{Union{ComplexF32, ComplexF64}}},
+ Tuple{Vararg{Union{ComplexF32, ComplexF64}}})
+tmerge_test(Tuple{ComplexF64, ComplexF64, ComplexF32}, Tuple{Vararg{Union{ComplexF32, ComplexF64}}},
+ Tuple{Vararg{Complex}}, false)
+tmerge_test(Tuple{}, Tuple{Complex, Vararg{Union{ComplexF32, ComplexF64}}},
+ Tuple{Vararg{Complex}})
+@test Core.Compiler.tmerge(Tuple{}, Union{Int16, Nothing, Tuple{ComplexF32, ComplexF32}}) ==
+ Union{Int16, Nothing, Tuple{Vararg{ComplexF32}}}
+@test Core.Compiler.tmerge(Int32, Union{Int16, Nothing, Tuple{ComplexF32, ComplexF32}}) ==
+ Union{Int16, Int32, Nothing, Tuple{ComplexF32, ComplexF32}}
+@test Core.Compiler.tmerge(Union{Int32, Nothing, Tuple{ComplexF32}}, Union{Int16, Nothing, Tuple{ComplexF32, ComplexF32}}) ==
+ Union{Int16, Int32, Nothing, Tuple{Vararg{ComplexF32}}}
 
 # issue 9770
 @noinline x9770() = false