[llvm] [SimplifyCFG] Fix use-after-free in switch-on-select remap when the compared value is a PHI (PR #223632)

[Akash Manna via llvm-commits]

https://github.com/akash-manna-sky created https://github.com/llvm/llvm-project/pull/223632

Fixes #223138

`simplifySwitchOnSelectRemap` turns `switch (select (X == C), K, X)` into `switch X` and retargets case `C` to wherever `K` goes. It dropped the old edge for `C` with `removePredecessor` before it installed `X` as the switch condition. When `X` is a PHI in that old destination, which is what happens when case `C` branches back to the switch's own block, dropping the edge leaves the PHI with a single input, so it gets folded and erased. The switch was then made to use the erased PHI. Everything after that was reading freed memory, which is why the reported assertion moves around between SimplifyCFG's dominator tree update, ScalarEvolution's `isSCEVable` check and instcombine depending on the run.

The fix is to set the switch condition to `X` before touching the CFG. If the PHI fold happens, its RAUW now rewrites the switch operand along with the compare and select, and the stale pointer is never used again.

>From 9a50a2b95aef05e7898b0ccb334c6f9504f9fc09 Mon Sep 17 00:00:00 2001
From: Akash Manna <akash.manna.mymail at gmail.com>
Date: Tue, 15 Sep 2026 14:25:11 +0530
Subject: [PATCH] [SimplifyCFG] Fix use-after-free in switch-on-select remap
 when the compared value is a PHI

simplifySwitchOnSelectRemap removed the old edge for case C before
installing X as the switch condition. When X is a PHI in that old
destination, e.g. a switch that loops back to its own block,
removePredecessor folds and erases the PHI, and the switch was then
pointed at the freed instruction.

Set the switch condition to X first so the RAUW performed by the PHI
fold updates the switch as well.

Fixes #223138
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     |  7 ++-
 .../SimplifyCFG/switch-select-remap.ll        | 46 +++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 6a2601487b178..668c0e819015a 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -5100,6 +5100,10 @@ bool SimplifyCFGOpt::simplifySwitchOnSelectRemap(SwitchInst *SI,
   BasicBlock *OldDest = CaseC->getCaseSuccessor();
   BasicBlock *BB = SI->getParent();
 
+  // Switch on X first: removePredecessor() below may fold X away if it is a
+  // PHI in OldDest, and the RAUW must update the switch condition too.
+  SI->setCondition(X);
+
   if (OldDest != DestFork) {
     if (!IsDefault)
       OldDest->removePredecessor(BB);
@@ -5176,8 +5180,7 @@ bool SimplifyCFGOpt::simplifySwitchOnSelectRemap(SwitchInst *SI,
     }
   }
 
-  // X replaces the condition so compare/select are now dead.
-  SI->setCondition(X);
+  // The compare/select are now dead.
   RecursivelyDeleteTriviallyDeadInstructions(Select);
   return true;
 }
diff --git a/llvm/test/Transforms/SimplifyCFG/switch-select-remap.ll b/llvm/test/Transforms/SimplifyCFG/switch-select-remap.ll
index 22bb3aaeca356..3bf4c2729c3e1 100644
--- a/llvm/test/Transforms/SimplifyCFG/switch-select-remap.ll
+++ b/llvm/test/Transforms/SimplifyCFG/switch-select-remap.ll
@@ -403,6 +403,52 @@ default:
   unreachable
 }
 
+; %x is a PHI in the switch's own block and the stale case for 4 branches back
+; to it. Retargeting that case removes the self-edge and folds the PHI away;
+; the switch used to be left pointing at the erased PHI (GH223138).
+define void @test_remap_retarget_self_loop_phi(i8 %a, i8 %b) {
+; CHECK-LABEL: define void @test_remap_retarget_self_loop_phi(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    switch i8 [[A]], label %[[BB1:.*]] [
+; CHECK-NEXT:      i8 4, label %[[BB2:.*]]
+; CHECK-NEXT:      i8 6, label %[[BB2]]
+; CHECK-NEXT:      i8 10, label %[[BB3:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    call void @func1()
+; CHECK-NEXT:    unreachable
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    call void @func2()
+; CHECK-NEXT:    unreachable
+; CHECK:       [[BB3]]:
+; CHECK-NEXT:    call void @func3()
+; CHECK-NEXT:    unreachable
+;
+entry:
+  br label %loop
+
+loop:
+  %x = phi i8 [ %a, %entry ], [ %b, %loop ]
+  %cmp = icmp eq i8 %x, 4
+  %key = select i1 %cmp, i8 6, i8 %x
+  switch i8 %key, label %bb1 [
+  i8 4, label %loop
+  i8 6, label %bb2
+  i8 10, label %bb3
+  ]
+
+bb1:
+  call void @func1()
+  unreachable
+bb2:
+  call void @func2()
+  unreachable
+bb3:
+  call void @func3()
+  unreachable
+}
+
 ; Negative test: %key (the select) is used by more than just the switch, so
 ; folding it away wouldn't actually remove the compare/select sequence -
 ; leave it alone.