From 3352d22abd66ae684d0c1ff5e08661ed303ce419 Mon Sep 17 00:00:00 2001
From: Cody Tapscott <topolarity@tapscott.me>
Date: Tue, 27 Jun 2023 14:01:31 -0400
Subject: [PATCH] optimize: Handle path-excluded `Core.ifelse` arguments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

It's possible for PiNodes to effectively imply statically the condition
of a Core.ifelse. For example:
```julia
    23 ─ %60  = Core.ifelse(%47, false, true)::Bool
    │    %61  = Core.ifelse(%47, %58, false)::Union{Missing, Bool}
    25 ─        goto #27 if not %60
    26 ─ %65  = π (%61, Bool)
    └───        ...
```

In basic block #26, the PiNode gives us enough information to conclude
that `%47 === false` if control flow reaches that point. The previous
code incorrectly assumed that this kind of pruning would only be done
for PhiNodes.

Resolves #50276.
---
 base/compiler/ssair/passes.jl | 17 +++++++++++-
 test/compiler/irpasses.jl     | 52 +++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
index f2ef2e9d47ee15..96656b0157780c 100644
--- a/base/compiler/ssair/passes.jl
+++ b/base/compiler/ssair/passes.jl
@@ -730,7 +730,7 @@ function perform_lifting!(compact::IncrementalCompact,
             new_node = Expr(:call, ifelse_func, condition) # Renamed then_result, else_result added below
             new_inst = NewInstruction(new_node, result_t, NoCallInfo(), old_inst[:line], old_inst[:flag])
 
-            ssa = insert_node!(compact, old_ssa, new_inst)
+            ssa = insert_node!(compact, old_ssa, new_inst, #= attach_after =# true)
             lifted_philikes[i] = LiftedPhilike(ssa, IfElseCall(new_node), true)
         end
         # lifting_cache[ckey] = ssa
@@ -767,6 +767,21 @@ function perform_lifting!(compact::IncrementalCompact,
             else_result = lifted_value(compact, old_node_ssa, else_result,
                                        lifted_philikes, lifted_leaves, reverse_mapping)
 
+            # In cases where the Core.ifelse condition is statically-known, e.g., thanks
+            # to a PiNode from a guarding conditional, replace with the remaining branch.
+            if (then_result === SKIP_TOKEN || else_result === SKIP_TOKEN)
+                only_result = (then_result === SKIP_TOKEN) ? else_result : then_result
+
+                # Replace Core.ifelse(%cond, %a, %b) with %a
+                compact[lf.ssa][:inst] = only_result
+                should_count && _count_added_node!(compact, only_result)
+
+                # Note: Core.ifelse(%cond, %a, %b) has observable effects (!throw), but since
+                # we have not deleted the preceding statement that this was copied from, we
+                # can delete it without changing whether those effects are observed.
+                continue
+            end
+
             @assert then_result !== SKIP_TOKEN && then_result !== UNDEF_TOKEN
             @assert else_result !== SKIP_TOKEN && else_result !== UNDEF_TOKEN
 
diff --git a/test/compiler/irpasses.jl b/test/compiler/irpasses.jl
index f27961c526559c..6db6fe837285a3 100644
--- a/test/compiler/irpasses.jl
+++ b/test/compiler/irpasses.jl
@@ -742,6 +742,58 @@ let m = Meta.@lower 1 + 1
     @test Core.Compiler.verify_ir(ir) === nothing
 end
 
+# A lifted Core.ifelse with an eliminated branch (#50276)
+let m = Meta.@lower 1 + 1
+    @assert Meta.isexpr(m, :thunk)
+    src = m.args[1]::CodeInfo
+    src.code = Any[
+        # block 1
+        #=  %1: =# Core.Argument(2),
+        # block 2
+        #=  %2: =# Expr(:call, Core.ifelse, SSAValue(1), true, missing),
+        #=  %3: =# GotoIfNot(SSAValue(2), 11),
+        # block 3
+        #=  %4: =# PiNode(SSAValue(2), Bool), # <-- This PiNode is the trigger of the bug, since it
+                                              #     means that only one branch of the Core.ifelse
+                                              #     is lifted.
+        #=  %5: =# GotoIfNot(false, 8),
+        # block 2
+        #=  %6: =# nothing,
+        #=  %7: =# GotoNode(8),
+        # block 4
+        #=  %8: =# PhiNode(Int32[5, 7], Any[SSAValue(4), SSAValue(6)]),
+        #               ^-- N.B. This PhiNode also needs to have a Union{ ... } type in order
+        #                   for lifting to be performed (it is skipped for e.g. `Bool`)
+        #
+        #=  %9: =# Expr(:call, isa, SSAValue(8), Missing),
+        #= %10: =# ReturnNode(SSAValue(9)),
+        # block 5
+        #= %11: =# ReturnNode(false),
+    ]
+    src.ssavaluetypes = Any[
+        Any,
+        Union{Missing, Bool},
+        Any,
+        Bool,
+        Any,
+        Missing,
+        Any,
+        Union{Nothing, Bool},
+        Bool,
+        Any,
+        Any
+    ]
+    nstmts = length(src.code)
+    src.codelocs = fill(one(Int32), nstmts)
+    src.ssaflags = fill(one(Int32), nstmts)
+    src.slotflags = fill(zero(UInt8), 3)
+    ir = Core.Compiler.inflate_ir(src)
+    Main.Base.IRShow.show_ir(stdout, ir)
+    @test Core.Compiler.verify_ir(ir) === nothing
+    ir = @test_nowarn Core.Compiler.sroa_pass!(ir)
+    @test Core.Compiler.verify_ir(ir) === nothing
+end
+
 # Issue #31546 - missing widenconst in SROA
 function f_31546(x)
     (a, b) = x == "r"  ? (false, false) :