From 94b7154a089b17bca235bc346a8b8d41f4f71e79 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Tue, 22 Mar 2022 11:44:08 +0900
Subject: [PATCH] inference: override `InterConditional` result with `Const`
 carefully (#44668)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

I found that a tricky thing can happen when constant inference derives
`Const`-result while non-constant inference has derived (non-constant)
`InterConditional` result beforehand. In such a case, currently we discard
the result with constant inference (since `!(Const ⊑ InterConditional)`),
but we can achieve more accuracy by not discarding that `Const`-information, e.g.:
```julia
julia> iszero_simple(x) = x === 0
iszero_simple (generic function with 1 method)

julia> @test Base.return_types() do
           iszero_simple(0) ? nothing : missing
       end |> only === Nothing
Test Passed
```
---
 base/compiler/abstractinterpretation.jl | 31 ++++++++++++++-----------
 test/compiler/inference.jl              |  7 ++++++
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl
index 2efb660a32fb5..3cc709bc35849 100644
--- a/base/compiler/abstractinterpretation.jl
+++ b/base/compiler/abstractinterpretation.jl
@@ -110,10 +110,9 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
             splitsigs = switchtupleunion(sig)
             for sig_n in splitsigs
                 result = abstract_call_method(interp, method, sig_n, svec(), multiple_matches, sv)
-                rt, edge = result.rt, result.edge
-                if edge !== nothing
-                    push!(edges, edge)
-                end
+                rt = result.rt
+                edge = result.edge
+                edge !== nothing && push!(edges, edge)
                 this_argtypes = isa(matches, MethodMatches) ? argtypes : matches.applicable_argtypes[i]
                 this_arginfo = ArgInfo(fargs, this_argtypes)
                 const_call_result = abstract_call_method_with_const_args(interp, result,
@@ -121,8 +120,10 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
                 effects = result.edge_effects
                 const_result = nothing
                 if const_call_result !== nothing
-                    if const_call_result.rt ⊑ rt
-                        (; rt, effects, const_result) = const_call_result
+                    const_rt = const_call_result.rt
+                    if const_rt ⊑ rt
+                        rt = const_rt
+                        (; effects, const_result) = const_call_result
                     end
                 end
                 tristate_merge!(sv, effects)
@@ -135,6 +136,8 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
                     break
                 end
             end
+            this_conditional = ignorelimited(this_rt)
+            this_rt = widenwrappedconditional(this_rt)
         else
             if infer_compilation_signature(interp)
                 # Also infer the compilation signature for this method, so it's available
@@ -151,10 +154,10 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
             end
 
             result = abstract_call_method(interp, method, sig, match.sparams, multiple_matches, sv)
-            this_rt, edge = result.rt, result.edge
-            if edge !== nothing
-                push!(edges, edge)
-            end
+            this_conditional = ignorelimited(result.rt)
+            this_rt = widenwrappedconditional(result.rt)
+            edge = result.edge
+            edge !== nothing && push!(edges, edge)
             # try constant propagation with argtypes for this match
             # this is in preparation for inlining, or improving the return result
             this_argtypes = isa(matches, MethodMatches) ? argtypes : matches.applicable_argtypes[i]
@@ -164,10 +167,12 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
             effects = result.edge_effects
             const_result = nothing
             if const_call_result !== nothing
-                this_const_rt = const_call_result.rt
-                # return type of const-prop' inference can be wider than  that of non const-prop' inference
+                this_const_conditional = ignorelimited(const_call_result.rt)
+                this_const_rt = widenwrappedconditional(const_call_result.rt)
+                # return type of const-prop' inference can be wider than that of non const-prop' inference
                 # e.g. in cases when there are cycles but cached result is still accurate
                 if this_const_rt ⊑ this_rt
+                    this_conditional = this_const_conditional
                     this_rt = this_const_rt
                     (; effects, const_result) = const_call_result
                 end
@@ -178,8 +183,6 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
                 any_const_result = true
             end
         end
-        this_conditional = ignorelimited(this_rt)
-        this_rt = widenwrappedconditional(this_rt)
         @assert !(this_conditional isa Conditional) "invalid lattice element returned from inter-procedural context"
         seen += 1
         rettype = tmerge(rettype, this_rt)
diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
index 218e484b2beca..e6f113f2b9062 100644
--- a/test/compiler/inference.jl
+++ b/test/compiler/inference.jl
@@ -2032,6 +2032,13 @@ end
         end
         @test ts == Any[Any]
     end
+
+    # a tricky case: if constant inference derives `Const` while non-constant inference has
+    # derived `InterConditional`, we should not discard that constant information
+    iszero_simple(x) = x === 0
+    @test Base.return_types() do
+        iszero_simple(0) ? nothing : missing
+    end |> only === Nothing
 end
 
 @testset "branching on conditional object" begin