From cd1880b8a11c82b620d5ff1254be6faf90d0449c Mon Sep 17 00:00:00 2001
From: Jameson Nash <vtjnash@gmail.com>
Date: Tue, 11 Jul 2023 15:52:25 -0400
Subject: [PATCH] Revert "SROA: generalize `unswitchtupleunion` optimization
 (#50502)"

This reverts commit 39952788ab337f1c85a9c51515de84020b665a74.
---
 base/compiler/ssair/passes.jl |  4 ++--
 base/compiler/typeutils.jl    | 45 +++++++++++++++--------------------
 test/compiler/irpasses.jl     | 31 ------------------------
 3 files changed, 21 insertions(+), 59 deletions(-)

diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
index 8f20ac28e3606..6e125d72d5e41 100644
--- a/base/compiler/ssair/passes.jl
+++ b/base/compiler/ssair/passes.jl
@@ -1107,8 +1107,8 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
         end
         struct_typ = widenconst(argextype(val, compact))
         struct_typ_unwrapped = unwrap_unionall(struct_typ)
-        if isa(struct_typ, Union)
-            struct_typ_unwrapped = unswitchtypeunion(struct_typ_unwrapped)
+        if isa(struct_typ, Union) && struct_typ <: Tuple
+            struct_typ_unwrapped = unswitchtupleunion(struct_typ_unwrapped)
         end
         if isa(struct_typ_unwrapped, Union) && is_isdefined
             lift_comparison!(isdefined, compact, idx, stmt, lifting_cache, 𝕃ₒ)
diff --git a/base/compiler/typeutils.jl b/base/compiler/typeutils.jl
index 7383ec2a440bf..f1794bb83c375 100644
--- a/base/compiler/typeutils.jl
+++ b/base/compiler/typeutils.jl
@@ -317,40 +317,33 @@ function unionall_depth(@nospecialize ua) # aka subtype_env_size
     return depth
 end
 
-# convert a Union of same `UnionAll` types to the `UnionAll` type whose parameter is the Unions
+# convert a Union of Tuple types to a Tuple of Unions
+unswitchtupleunion(u::Union) = unswitchtypeunion(u, Tuple.name)
+
 function unswitchtypeunion(u::Union, typename::Union{Nothing,Core.TypeName}=nothing)
     ts = uniontypes(u)
     n = -1
     for t in ts
-        t isa DataType || return u
-        if typename === nothing
-            typename = t.name
-        elseif typename !== t.name
-            return u
-        end
-        params = t.parameters
-        np = length(params)
-        if np == 0 || isvarargtype(params[end])
-            return u
-        end
-        if n == -1
-            n = np
-        elseif n ≠ np
+        if t isa DataType
+            if typename === nothing
+                typename = t.name
+            elseif typename !== t.name
+                return u
+            end
+            if length(t.parameters) != 0 && !isvarargtype(t.parameters[end])
+                if n == -1
+                    n = length(t.parameters)
+                elseif n != length(t.parameters)
+                    return u
+                end
+            end
+        else
             return u
         end
     end
     Head = (typename::Core.TypeName).wrapper
-    hparams = Any[]
-    for i = 1:n
-        uparams = Any[]
-        for t in ts
-            tpᵢ = (t::DataType).parameters[i]
-            tpᵢ isa Type || return u
-            push!(uparams, tpᵢ)
-        end
-        push!(hparams, Union{uparams...})
-    end
-    return Head{hparams...}
+    unionparams = Any[ Union{Any[(t::DataType).parameters[i] for t in ts]...} for i in 1:n ]
+    return Head{unionparams...}
 end
 
 function unwraptv_ub(@nospecialize t)
diff --git a/test/compiler/irpasses.jl b/test/compiler/irpasses.jl
index c163d141fc08f..1cba2d2ee0006 100644
--- a/test/compiler/irpasses.jl
+++ b/test/compiler/irpasses.jl
@@ -1390,34 +1390,3 @@ function wrap1_wrap1_wrapper(b, x, y)
 end
 @test wrap1_wrap1_wrapper(true, 1, 1.0) === 1.0
 @test wrap1_wrap1_wrapper(false, 1, 1.0) === 1
-
-# Test unswitching-union optimization within SRO Apass
-function sroaunswitchuniontuple(c, x1, x2)
-    t = c ? (x1,) : (x2,)
-    return getfield(t, 1)
-end
-struct SROAUnswitchUnion1{T}
-    x::T
-end
-struct SROAUnswitchUnion2{S,T}
-    x::T
-    @inline SROAUnswitchUnion2{S}(x::T) where {S,T} = new{S,T}(x)
-end
-function sroaunswitchunionstruct1(c, x1, x2)
-    x = c ? SROAUnswitchUnion1(x1) : SROAUnswitchUnion1(x2)
-    return getfield(x, :x)
-end
-function sroaunswitchunionstruct2(c, x1, x2)
-    x = c ? SROAUnswitchUnion2{:a}(x1) : SROAUnswitchUnion2{:a}(x2)
-    return getfield(x, :x)
-end
-let src = code_typed1(sroaunswitchuniontuple, Tuple{Bool, Int, Float64})
-    @test count(isnew, src.code) == 0
-    @test count(iscall((src, getfield)), src.code) == 0
-end
-let src = code_typed1(sroaunswitchunionstruct1, Tuple{Bool, Int, Float64})
-    @test count(isnew, src.code) == 0
-    @test count(iscall((src, getfield)), src.code) == 0
-end
-@test sroaunswitchunionstruct2(true, 1, 1.0) === 1
-@test sroaunswitchunionstruct2(false, 1, 1.0) === 1.0