inference: stronger state update from branching

This change addresses <JuliaLang#40832 (review)>
aviatesk · Jun 3, 2021 · ae85281 · ae85281
1 parent 8c66f95
commit ae85281
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git 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
@@ -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,13 +1843,24 @@ 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
         else
             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)