Fix typeassert dropping PartialStruct information (#36622)

There was a complaint on discourse that Generator splats weren't
infering precisely enough [1]. I thought the recent improvement
to abstract iteration precision would have fixed that, but that
turns out not to be the case of a stupid reason. At the moment,
any type assert (even something silly like `x::Any`), will drop
`PartialStruct` information. Now, generators have an
`::Tuple{Any, Any}` typeassert, which prevented the new iteration
logic from working properly. This fixes typasserts, to not drop
this information. The first couple of conditions that just
return the original `PartialStruct` if the typeassert has no
effect are doing most of the work here. As an add-on, I also
threw in the logic to narrow the PartialStruct if the typeassert
actually introduces additional information. It's not a particularly
common code path, but might as well.

[1] https://discourse.julialang.org/t/tuples-iterators-splats-and-type-stability/42849

Author: Keno

base/compiler/tfuncs.jl

@@ -485,18 +485,38 @@ add_tfunc(typeof, 1, 1, typeof_tfunc, 0)
 function typeassert_tfunc(@nospecialize(v), @nospecialize(t))
     t = instanceof_tfunc(t)[1]
     t === Any && return v
-    if isa(v, Const)
-        if !has_free_typevars(t) && !isa(v.val, t)
-            return Bottom
-        end
-        return v
-    elseif isa(v, Conditional)
-        if !(Bool <: t)
-            return Bottom
+    function typeassert_type_instance(@nospecialize(v), @nospecialize(t))
+        if isa(v, Const)
+            if !has_free_typevars(t) && !isa(v.val, t)
+                return Bottom
+            end
+            return v
+        elseif isa(v, PartialStruct)
+            has_free_typevars(t) && return v
+            widev = widenconst(v)
+            if widev <: t
+                return v
+            elseif typeintersect(widev, t) === Bottom
+                return Bottom
+            end
+            @assert widev <: Tuple
+            new_fields = Vector{Any}(undef, length(v.fields))
+            for i = 1:length(new_fields)
+                new_fields[i] = typeassert_type_instance(v.fields[i], getfield_tfunc(t, Const(i)))
+                if new_fields[i] === Bottom
+                    return Bottom
+                end
+            end
+            return tuple_tfunc(new_fields)
+        elseif isa(v, Conditional)
+            if !(Bool <: t)
+                return Bottom
+            end
+            return v
         end
-        return v
+        return typeintersect(widenconst(v), t)
     end
-    return typeintersect(widenconst(v), t)
+    return typeassert_type_instance(v, t)
 end
 add_tfunc(typeassert, 2, 2, typeassert_tfunc, 4)
 

test/compiler/inference.jl

@@ -2716,3 +2716,21 @@ function f_apply_union_split(fs, x)
 end
 
 @test Base.return_types(f_apply_union_split, Tuple{Tuple{typeof(sqrt), typeof(abs)}, Int64}) == Any[Union{Int64, Float64}]
+
+# Precision of typeassert with PartialStruct
+function f_typ_assert(x::Int)
+    y = (x, 1)
+    y = y::Any
+    Val{y[2]}
+end
+@test Base.return_types(f_typ_assert, (Int,)) == Any[Type{Val{1}}]
+
+function f_typ_assert2(x::Any)
+    y = (x::Union{Int, Float64}, 1)
+    y = y::Tuple{Int, Any}
+    (y[1], Val{y[2]}())
+end
+@test Base.return_types(f_typ_assert2, (Any,)) == Any[Tuple{Int, Val{1}}]
+
+f_generator_splat(t::Tuple) = tuple((identity(l) for l in t)...)
+@test Base.return_types(f_generator_splat, (Tuple{Symbol, Int64, Float64},)) == Any[Tuple{Symbol, Int64, Float64}]