[llvm] [SimplifyCFG] Handle switch-to-select remapping correctly (PR #213302)

[Yingwei Zheng via llvm-commits]

================
@@ -4999,8 +4999,70 @@ bool SimplifyCFGOpt::simplifyTerminatorOnSelect(Instruction *OldTerm,
 //   (switch (select cond, X, Y)) on constant X, Y
 // with a branch - conditional if X and Y lead to distinct BBs,
 // unconditional otherwise.
+//
+// Also folds switch(select(icmp eq X, C, K, X)) into switch(X), retargeting
+// (or adding) the case for C to wherever K currently dispatches to:
+//   %cmp = icmp eq T %x, C
+//   %key = select i1 %cmp, T K, T %x
+//   switch T %key, label %default [ T K, label %case_k ... ]
+// becomes
+//   switch T %x, label %default [ T C, label %case_k
+//                                  T K, label %case_k ... ]
 bool SimplifyCFGOpt::simplifySwitchOnSelect(SwitchInst *SI,
                                             SelectInst *Select) {
+  CmpPredicate Pred;
+  Value *X;
+  ConstantInt *C;
+  if (Select->hasOneUse() &&
+      match(Select->getCondition(),
+            m_OneUse(m_c_ICmp(Pred, m_Value(X), m_ConstantInt(C)))) &&
+      ICmpInst::isEquality(Pred)) {
+    // Normalize to select(icmp eq X, C, K, X).
+    Value *TrueVal = Select->getTrueValue();
+    Value *FalseVal = Select->getFalseValue();
+    if (Pred == ICmpInst::ICMP_NE)
+      std::swap(TrueVal, FalseVal);
+    if (auto *K = (FalseVal == X) ? dyn_cast<ConstantInt>(TrueVal) : nullptr) {
+      BasicBlock *DestFork = SI->findCaseValue(K)->getCaseSuccessor();
+      auto CaseC = SI->findCaseValue(C);
+      bool IsDefault = CaseC == SI->case_default();
+      // Save before setSuccessor()/addCase() change it.
+      BasicBlock *OldDest = CaseC->getCaseSuccessor();
+      BasicBlock *BB = SI->getParent();
+
+      if (OldDest != DestFork) {
+        // Case list is changing so we should drop stale profile weights.
+        SI->setMetadata(LLVMContext::MD_prof, nullptr);
+        if (IsDefault)
+          SI->addCase(C, DestFork);
+        else
+          CaseC->setSuccessor(DestFork);
+        // Not a new edge (BB->DestFork exists via K), just adding the PHI
+        // entry.
+        addPredecessorToBlock(DestFork, BB, BB);
+
+        if (!IsDefault) {
+          // Edge to OldDest is gone only if nothing else still uses it.
+          bool OldDestStillTargeted =
+              SI->getDefaultDest() == OldDest ||
+              any_of(SI->cases(), [&](const SwitchInst::CaseHandle &Case) {
+                return Case.getCaseIndex() != CaseC->getCaseIndex() &&
+                       Case.getCaseSuccessor() == OldDest;
+              });
+          for (PHINode &PN : make_early_inc_range(OldDest->phis()))
----------------
dtcxzyw wrote:

Use `OldDest->removePredecessor` instead.


https://github.com/llvm/llvm-project/pull/213302