[SimplifyCFG] Avoid use-after-free when removing incoming values from PHI nodes

[dtcxzyw]

PHINode::removeIncomingValue removes itself when there are no incoming edges. Then we cannot use it to retrieve the next instruction.

Closes #165301.

Note: I found that other places contain the same pattern. Should we make DeletePHIIfEmpty default to false?

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

PHINode::removeIncomingValue removes itself when there are no incoming edges. Then we cannot use it to retrieve the next instruction.

Closes #165301.

Note: I found that other places contain the same pattern. Should we make DeletePHIIfEmpty default to false?

---

Full diff: https://github.com/llvm/llvm-project/pull/165744.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+4-4)
• (added) llvm/test/Transforms/SimplifyCFG/pr165301.ll (+26)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index b03fb6213d61c..7f6d779687e94 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -5977,14 +5977,14 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
   }
 
   // Prune obsolete incoming values off the successors' PHI nodes.
-  for (auto BBI = Dest->begin(); isa<PHINode>(BBI); ++BBI) {
+  for (auto &PHI : make_early_inc_range(Dest->phis())) {
     unsigned PreviousEdges = Cases->size();
     if (Dest == SI->getDefaultDest())
       ++PreviousEdges;
     for (unsigned I = 0, E = PreviousEdges - 1; I != E; ++I)
-      cast<PHINode>(BBI)->removeIncomingValue(SI->getParent());
+      PHI.removeIncomingValue(SI->getParent());
   }
-  for (auto BBI = OtherDest->begin(); isa<PHINode>(BBI); ++BBI) {
+  for (auto &PHI : make_early_inc_range(OtherDest->phis())) {
     unsigned PreviousEdges = OtherCases->size();
     if (OtherDest == SI->getDefaultDest())
       ++PreviousEdges;
@@ -5993,7 +5993,7 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
     if (NewBI->isUnconditional())
       ++E;
     for (unsigned I = 0; I != E; ++I)
-      cast<PHINode>(BBI)->removeIncomingValue(SI->getParent());
+      PHI.removeIncomingValue(SI->getParent());
   }
 
   // Clean up the default block - it may have phis or other instructions before
diff --git a/llvm/test/Transforms/SimplifyCFG/pr165301.ll b/llvm/test/Transforms/SimplifyCFG/pr165301.ll
new file mode 100644
index 0000000000000..4a539d77af3cb
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/pr165301.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes="simplifycfg<switch-range-to-icmp>" < %s | FileCheck %s
+
+; Make sure there's no use after free when removing incoming values from PHI nodes
+
+define i32 @pr165301(i1 %cond) {
+; CHECK-LABEL: define i32 @pr165301(
+; CHECK-SAME: i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[SWITCHBB:.*]]
+; CHECK:       [[SWITCHBB]]:
+; CHECK-NEXT:    br label %[[SWITCHBB]]
+;
+entry:
+  br label %switchbb
+
+switchbb:
+  switch i1 %cond, label %default [
+  i1 false, label %switchbb
+  i1 true, label %switchbb
+  ]
+
+default:
+  %phi.lcssa = phi i32 [ 0, %switchbb ]
+  ret i32 %phi.lcssa
+}

[nikic]

LGTM