inference: apply tmerge limit elementwise to the Union (JuliaLang#50927)

This allows forming larger unions, as long as each element in the Union
is both relatively distinct and relatively simple. For example:

    tmerge(Base.BitSigned, Nothing) == Union{Nothing, Int128, Int16, Int32, Int64, Int8}
    tmerge(Tuple{Base.BitSigned, Int}, Nothing) == Union{Nothing, Tuple{Any, Int64}}
    tmerge(AbstractVector{Int}, Vector) == AbstractVector

Disables a test from dc8d885, which does not seem possible to handle currently.

This makes somewhat drastic changes to make this algorithm more
commutative and simpler, since we dropped the final widening to `Any`.

Co-authored-by: pchintalapudi <34727397+pchintalapudi@users.noreply.github.com>
Co-authored-by: Oscar Smith <oscardssmith@gmail.com>

Author: 3 people

base/compiler/abstractinterpretation.jl

@@ -3122,11 +3122,7 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
                 ssavaluetypes[currpc] = Any
                 continue
             end
-            if !isempty(frame.ssavalue_uses[currpc])
-                record_ssa_assign!(𝕃ᵢ, currpc, type, frame)
-            else
-                ssavaluetypes[currpc] = type
-            end
+            record_ssa_assign!(𝕃ᵢ, currpc, type, frame)
         end # for currpc in bbstart:bbend
 
         # Case 1: Fallthrough termination

base/compiler/inferencestate.jl

@@ -576,11 +576,8 @@ _topmod(sv::InferenceState) = _topmod(frame_module(sv))
 function record_ssa_assign!(𝕃ᵢ::AbstractLattice, ssa_id::Int, @nospecialize(new), frame::InferenceState)
     ssavaluetypes = frame.ssavaluetypes
     old = ssavaluetypes[ssa_id]
-    if old === NOT_FOUND || !⊑(𝕃ᵢ, new, old)
-        # typically, we expect that old ⊑ new (that output information only
-        # gets less precise with worse input information), but to actually
-        # guarantee convergence we need to use tmerge here to ensure that is true
-        ssavaluetypes[ssa_id] = old === NOT_FOUND ? new : tmerge(𝕃ᵢ, old, new)
+    if old === NOT_FOUND || !is_lattice_equal(𝕃ᵢ, new, old)
+        ssavaluetypes[ssa_id] = new
         W = frame.ip
         for r in frame.ssavalue_uses[ssa_id]
             if was_reached(frame, r)

base/compiler/typelimits.jl

@@ -292,10 +292,17 @@ union_count_abstract(x::Union) = union_count_abstract(x.a) + union_count_abstrac
 union_count_abstract(@nospecialize(x)) = !isdispatchelem(x)
 
 function issimpleenoughtype(@nospecialize t)
+    ut = unwrap_unionall(t)
+    ut isa DataType && ut.name.wrapper == t && return true
     return unionlen(t) + union_count_abstract(t) <= MAX_TYPEUNION_LENGTH &&
            unioncomplexity(t) <= MAX_TYPEUNION_COMPLEXITY
 end
 
+# We may want to apply a stricter limit than issimpleenoughtype to
+# tupleelements individually, to try to keep the whole tuple under the limit,
+# even after complicated recursion and other operations on it elsewhere
+const issimpleenoughtupleelem = issimpleenoughtype
+
 # A simplified type_more_complex query over the extended lattice
 # (assumes typeb ⊑ typea)
 @nospecializeinfer function issimplertype(𝕃::AbstractLattice, @nospecialize(typea), @nospecialize(typeb))
@@ -679,6 +686,33 @@ end
     return tmerge_types_slow(typea, typeb)
 end
 
+@nospecializeinfer @noinline function tname_intersect(aname::Core.TypeName, bname::Core.TypeName)
+    aname === bname && return aname
+    if !isabstracttype(aname.wrapper) && !isabstracttype(bname.wrapper)
+        return nothing # fast path
+    end
+    Any.name === aname && return aname
+    a = unwrap_unionall(aname.wrapper)
+    heighta = 0
+    while a !== Any
+        heighta += 1
+        a = a.super
+    end
+    b = unwrap_unionall(bname.wrapper)
+    heightb = 0
+    while b !== Any
+        b.name === aname && return aname
+        heightb += 1
+        b = b.super
+    end
+    a = unwrap_unionall(aname.wrapper)
+    while heighta > heightb
+        a = a.super
+        heighta -= 1
+    end
+    return a.name === bname ? bname : nothing
+end
+
 @nospecializeinfer @noinline function tmerge_types_slow(@nospecialize(typea::Type), @nospecialize(typeb::Type))
     # collect the list of types from past tmerge calls returning Union
     # and then reduce over that list
@@ -702,74 +736,95 @@ end
     # see if any of the union elements have the same TypeName
     # in which case, simplify this tmerge by replacing it with
     # the widest possible version of itself (the wrapper)
+    simplify = falses(length(types))
     for i in 1:length(types)
+        typenames[i] === Any.name && continue
         ti = types[i]
         for j in (i + 1):length(types)
-            if typenames[i] === typenames[j]
+            typenames[j] === Any.name && continue
+            ijname = tname_intersect(typenames[i], typenames[j])
+            if !(ijname === nothing)
                 tj = types[j]
                 if ti <: tj
                     types[i] = Union{}
                     typenames[i] = Any.name
+                    simplify[i] = false
+                    simplify[j] = true
                     break
                 elseif tj <: ti
                     types[j] = Union{}
                     typenames[j] = Any.name
+                    simplify[j] = false
+                    simplify[i] = true
                 else
-                    if typenames[i] === Tuple.name
+                    if ijname === 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,
                         widen = tuplemerge(unwrap_unionall(ti)::DataType, unwrap_unionall(tj)::DataType)
                         widen = rewrap_unionall(rewrap_unionall(widen, ti), tj)
+                        simplify[j] = false
                     else
-                        wr = typenames[i].wrapper
+                        wr = ijname.wrapper
                         uw = unwrap_unionall(wr)::DataType
                         ui = unwrap_unionall(ti)::DataType
+                        while ui.name !== ijname
+                            ui = ui.super
+                        end
                         uj = unwrap_unionall(tj)::DataType
-                        merged = wr
+                        while uj.name !== ijname
+                            uj = uj.super
+                        end
+                        p = Vector{Any}(undef, length(uw.parameters))
+                        usep = true
+                        widen = wr
                         for k = 1:length(uw.parameters)
                             ui_k = ui.parameters[k]
                             if ui_k === uj.parameters[k] && !has_free_typevars(ui_k)
-                                merged = merged{ui_k}
+                                p[k] = ui_k
+                                usep = true
                             else
-                                merged = merged{uw.parameters[k]}
+                                p[k] = uw.parameters[k]
                             end
                         end
-                        widen = rewrap_unionall(merged, wr)
+                        if usep
+                            widen = rewrap_unionall(wr{p...}, wr)
+                        end
+                        simplify[j] = !usep
                     end
                     types[i] = Union{}
                     typenames[i] = Any.name
+                    simplify[i] = false
                     types[j] = widen
                     break
                 end
             end
         end
     end
-    u = Union{types...}
-    # don't let type unions get too big, if the above didn't reduce it enough
-    if issimpleenoughtype(u)
-        return u
-    end
-    # don't let the slow widening of Tuple cause the whole type to grow too fast
+    # don't let elements of the union get too big, if the above didn't reduce something enough
     # Specifically widen Tuple{..., Union{lots of stuff}...} to Tuple{..., Any, ...}
+    # Don't let Val{<:Val{<:Val}} keep nesting abstract levels either
     for i in 1:length(types)
+        simplify[i] || continue
+        ti = types[i]
+        issimpleenoughtype(ti) && continue
         if typenames[i] === Tuple.name
-            ti = types[i]
-            tip = (unwrap_unionall(types[i])::DataType).parameters
+            # otherwise we need to do a simple version of tuplemerge for one element now
+            tip = (unwrap_unionall(ti)::DataType).parameters
             lt = length(tip)
             p = Vector{Any}(undef, lt)
             for j = 1:lt
                 ui = tip[j]
-                p[j] = (unioncomplexity(ui)==0) ? ui : isvarargtype(ui) ? Vararg : Any
+                p[j] = issimpleenoughtupleelem(unwrapva(ui)) ? ui : isvarargtype(ui) ? Vararg : Any
             end
             types[i] = rewrap_unionall(Tuple{p...}, ti)
+        else
+            # this element is not simple enough yet, make it so now
+            types[i] = typenames[i].wrapper
         end
     end
     u = Union{types...}
-    if issimpleenoughtype(u)
-        return u
-    end
-    return Any
+    return u
 end
 
 # the inverse of switchtupleunion, with limits on max element union size
@@ -791,7 +846,7 @@ function tuplemerge(a::DataType, b::DataType)
     p = Vector{Any}(undef, lt + vt)
     for i = 1:lt
         ui = Union{ap[i], bp[i]}
-        p[i] = issimpleenoughtype(ui) ? ui : Any
+        p[i] = issimpleenoughtupleelem(ui) ? ui : Any
     end
     # merge the remaining tail into a single, simple Tuple{Vararg{T}} (#22120)
     if vt
@@ -809,8 +864,10 @@ function tuplemerge(a::DataType, b::DataType)
                 #   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
+                #   c.f. tname_intersect in the algorithm above
+                hasfree = has_free_typevars(ti)
+                if hasfree || !(ti <: tail)
+                    if !hasfree && tail <: ti
                         tail = ti # widen to ti
                     else
                         uw = unwrap_unionall(tail)
@@ -838,11 +895,16 @@ function tuplemerge(a::DataType, b::DataType)
                         end
                     end
                 end
-                tail === Any && return Tuple # short-circuit loop
+                tail === Any && return Tuple # short-circuit loops
             end
         end
         @assert !(tail === Union{})
-        p[lt + 1] = Vararg{tail}
+        if !issimpleenoughtupleelem(tail) || tail === Any
+            p[lt + 1] = Vararg
+            lt == 0 && return Tuple
+        else
+            p[lt + 1] = Vararg{tail}
+        end
     end
     return Tuple{p...}
 end

base/compiler/typeutils.jl

@@ -291,7 +291,7 @@ end
 unioncomplexity(@nospecialize x) = _unioncomplexity(x)::Int
 function _unioncomplexity(@nospecialize x)
     if isa(x, DataType)
-        x.name === Tuple.name || isvarargtype(x) || return 0
+        x.name === Tuple.name || return 0
         c = 0
         for ti in x.parameters
             c = max(c, unioncomplexity(ti))
@@ -302,7 +302,7 @@ function _unioncomplexity(@nospecialize x)
     elseif isa(x, UnionAll)
         return max(unioncomplexity(x.body), unioncomplexity(x.var.ub))
     elseif isa(x, TypeofVararg)
-        return isdefined(x, :T) ? unioncomplexity(x.T) : 0
+        return isdefined(x, :T) ? unioncomplexity(x.T) + 1 : 1
     else
         return 0
     end