From ae85281ce8f3594124665aba6462c22f398c8d11 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Sun, 23 May 2021 13:30:28 +0900
Subject: [PATCH] inference: stronger state update from branching

This change addresses
<https://github.com/JuliaLang/julia/pull/40832#pullrequestreview-663397013>
---
 base/compiler/abstractinterpretation.jl | 11 ++++++-----
 test/compiler/inference.jl              | 14 +++++++++++++-
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl
index e5607a6047ab1e..a9ac65cca9f08e 100644
--- a/base/compiler/abstractinterpretation.jl
+++ b/base/compiler/abstractinterpretation.jl
@@ -1734,11 +1734,6 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
                     empty!(frame.pclimitations)
                     break
                 end
-                if !(isa(condt, Const) || isa(condt, Conditional)) && isa(condx, SlotNumber)
-                    # if this non-`Conditional` object is a slot, we form and propagate
-                    # the conditional constraint on it
-                    condt = Conditional(condx, Const(true), Const(false))
-                end
                 condval = maybe_extract_const_bool(condt)
                 l = stmt.dest::Int
                 if !isempty(frame.pclimitations)
@@ -1760,6 +1755,12 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
                         changes_else = conditional_changes(changes_else, condt.elsetype, condt.var)
                         changes      = conditional_changes(changes,      condt.vtype,    condt.var)
                     end
+                    if isa(condx, SlotNumber)
+                        tfalse = isa(condt, Conditional) ? Conditional(condt.var, Bottom, condt.elsetype) : Const(false)
+                        ttrue  = isa(condt, Conditional) ? Conditional(condt.var, condt.vtype, Bottom)    : Const(true)
+                        changes_else = add_state_change!(changes_else, condx, tfalse, true)
+                        changes      = add_state_change!(changes,      condx, ttrue,  true)
+                    end
                     newstate_else = stupdate!(states[l], changes_else)
                     if newstate_else !== nothing
                         # add else branch to active IP list
diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
index 34e3db8c989153..b622afa70772c0 100644
--- a/test/compiler/inference.jl
+++ b/test/compiler/inference.jl
@@ -1833,7 +1833,8 @@ end
     end == Any[Tuple{Int,Int}]
 end
 
-@testset "conditional constraint propagation from non-`Conditional` object" begin
+@testset "state update on branching" begin
+    # refine condition type into constant boolean value on branching
     @test Base.return_types((Bool,)) do b
         if b
             return !b ? nothing : 1 # ::Int
@@ -1842,6 +1843,7 @@ end
         end
     end == Any[Int]
 
+    # even when the original type isn't boolean type
     @test Base.return_types((Any,)) do b
         if b
             return b # ::Bool
@@ -1849,6 +1851,16 @@ end
             return nothing
         end
     end == Any[Union{Bool,Nothing}]
+
+    # even when the original type is `Conditional`
+    @test Base.return_types((Any,)) do a
+        b = isa(a, Int)
+        if b
+            return !b ? nothing : a # ::Int
+        else
+            return 0
+        end
+    end == Any[Int]
 end
 
 function f25579(g)