[llvm] [FixIrreducible][UnifyLoopExits] Fix callbr multiedge splitting (PR #207598)

[Robert Imschweiler via llvm-commits]

================
@@ -726,42 +726,62 @@ static bool updateCycleLoopInfo(TI *LCI, BasicBlock *CallBrBlock,
 }
 
 BasicBlock *llvm::SplitCallBrEdge(BasicBlock *CallBrBlock, BasicBlock *Succ,
-                                  unsigned SuccIdx, DomTreeUpdater *DTU,
-                                  CycleInfo *CI, LoopInfo *LI,
-                                  bool *UpdatedLI) {
+                                  unsigned SuccIdx, BasicBlock *CallBrTarget,
+                                  DomTreeUpdater *DTU, CycleInfo *CI,
+                                  LoopInfo *LI, bool *UpdatedLI) {
   CallBrInst *CallBr = dyn_cast<CallBrInst>(CallBrBlock->getTerminator());
   assert(CallBr && "expected callbr terminator");
   assert(SuccIdx < CallBr->getNumSuccessors() &&
          Succ == CallBr->getSuccessor(SuccIdx) && "invalid successor index");
 
+  bool ReusesCallBrTarget = CallBrTarget;
   // Create a new block between callbr and the specified successor.
   // splitBlockBefore cannot be re-used here since it cannot split if the split
   // point is a PHI node (because BasicBlock::splitBasicBlockBefore cannot
   // handle that). But we don't need to rewire every part of a potential PHI
   // node. We only care about the edge between CallBrBlock and the original
   // successor.
-  BasicBlock *CallBrTarget =
-      BasicBlock::Create(CallBrBlock->getContext(),
-                         CallBrBlock->getName() + ".target." + Succ->getName(),
-                         CallBrBlock->getParent());
-  // Rewire control flow from the new target block to the original successor.
-  Succ->replacePhiUsesWith(CallBrBlock, CallBrTarget);
+  if (!ReusesCallBrTarget) {
+    CallBrTarget = BasicBlock::Create(CallBrBlock->getContext(),
+                                      CallBrBlock->getName() + ".target." +
+                                          Succ->getName(),
+                                      CallBrBlock->getParent());
+    // Jump from the new target block to the original successor.
+    UncondBrInst::Create(Succ, CallBrTarget);
+    // Replace a single incoming value with the callbr target block. We cannot
+    // use replacePhiUsesWith, as this would replace the value for every edge
+    // from the callbr block to succ.
+    for (PHINode &PN : Succ->phis()) {
+      int BBIdx = PN.getBasicBlockIndex(CallBrBlock);
+      assert(BBIdx != -1 && "expected incoming value form callbr block");
+      PN.setIncomingBlock(BBIdx, CallBrTarget);
+    }
+  } else {
+    for (PHINode &PN : Succ->phis()) {
----------------
ro-i wrote:

no braces needed

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