From caf97c36ab01297361a24926f7371de1a358300f Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Fri, 8 Mar 2024 01:15:35 +0900
Subject: [PATCH] inlining: use method match signature for union-spliting
 (#53600)

In cases where the results of constant inference, like concrete-eval,
are used for union-split inlining, `isa`-blocks are generated using the
`result.edge.specTypes` stored within each `result`. However, it's been
found that the `edge` returned by abstract interpretation may have been
widened by the new `@nospecializeinfer`, which can result in invalid
union-splitting. To address this problem, this commit tweaks the
inlining algorithm so that it performs union-split inlining using the
original signatures that abstract interpretation used for union-split
inference, by using `match::MethodMatch`.

- fixes #53590
---
 base/compiler/ssair/inlining.jl | 23 ++++++++++++-----------
 test/compiler/inference.jl      | 12 ++++++++++++
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/base/compiler/ssair/inlining.jl b/base/compiler/ssair/inlining.jl
index a89583494ee137..3e57bc8503e06b 100644
--- a/base/compiler/ssair/inlining.jl
+++ b/base/compiler/ssair/inlining.jl
@@ -1331,11 +1331,11 @@ function handle_any_const_result!(cases::Vector{InliningCase},
     @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
     allow_abstract::Bool, allow_typevars::Bool)
     if isa(result, ConcreteResult)
-        return handle_concrete_result!(cases, result, info, state)
+        return handle_concrete_result!(cases, result, match, info, state)
     elseif isa(result, SemiConcreteResult)
-        return handle_semi_concrete_result!(cases, result, info, flag, state; allow_abstract)
+        return handle_semi_concrete_result!(cases, result, match, info, flag, state; allow_abstract)
     elseif isa(result, ConstPropResult)
-        return handle_const_prop_result!(cases, result, info, flag, state; allow_abstract, allow_typevars)
+        return handle_const_prop_result!(cases, result, match, info, flag, state; allow_abstract, allow_typevars)
     else
         @assert result === nothing || result isa VolatileInferenceResult
         return handle_match!(cases, match, argtypes, info, flag, state; allow_abstract, allow_typevars, volatile_inf_result = result)
@@ -1481,11 +1481,11 @@ function handle_match!(cases::Vector{InliningCase},
     return true
 end
 
-function handle_const_prop_result!(cases::Vector{InliningCase},
-    result::ConstPropResult, @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
+function handle_const_prop_result!(cases::Vector{InliningCase}, result::ConstPropResult,
+    match::MethodMatch, @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
     allow_abstract::Bool, allow_typevars::Bool)
     mi = result.result.linfo
-    spec_types = mi.specTypes
+    spec_types = match.spec_types
     allow_abstract || isdispatchtuple(spec_types) || return false
     if !validate_sparams(mi.sparam_vals)
         (allow_typevars && !may_have_fcalls(mi.def::Method)) || return false
@@ -1520,10 +1520,10 @@ function semiconcrete_result_item(result::SemiConcreteResult,
 end
 
 function handle_semi_concrete_result!(cases::Vector{InliningCase}, result::SemiConcreteResult,
-        @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
-        allow_abstract::Bool)
+    match::MethodMatch, @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
+    allow_abstract::Bool)
     mi = result.mi
-    spec_types = mi.specTypes
+    spec_types = match.spec_types
     allow_abstract || isdispatchtuple(spec_types) || return false
     validate_sparams(mi.sparam_vals) || return false
     item = semiconcrete_result_item(result, info, flag, state)
@@ -1532,10 +1532,11 @@ function handle_semi_concrete_result!(cases::Vector{InliningCase}, result::SemiC
     return true
 end
 
-function handle_concrete_result!(cases::Vector{InliningCase}, result::ConcreteResult, @nospecialize(info::CallInfo), state::InliningState)
+function handle_concrete_result!(cases::Vector{InliningCase}, result::ConcreteResult,
+    match::MethodMatch, @nospecialize(info::CallInfo), state::InliningState)
     case = concrete_result_item(result, info, state)
     case === nothing && return false
-    push!(cases, InliningCase(result.mi.specTypes, case))
+    push!(cases, InliningCase(match.spec_types, case))
     return true
 end
 
diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
index a9b1d89d40e9d4..d6cd231dfdd588 100644
--- a/test/compiler/inference.jl
+++ b/test/compiler/inference.jl
@@ -5627,3 +5627,15 @@ end
 
 # Issue #52613
 @test (code_typed((Any,)) do x; TypeVar(x...); end)[1][2] === TypeVar
+
+# https://github.com/JuliaLang/julia/issues/53590
+func53590(b) = b ? Int : Float64
+function issue53590(b1, b2)
+    T1 = func53590(b1)
+    T2 = func53590(b2)
+    return typejoin(T1, T2)
+end
+@test issue53590(true, true) == Int
+@test issue53590(true, false) == Real
+@test issue53590(false, false) == Float64
+@test issue53590(false, true) == Real