From 8201e3ef5c84561260218bc041209611aac690e3 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Mon, 18 Jul 2022 12:25:26 +0200
Subject: [PATCH] [BasicBlockUtils] Don't drop callbr with unique successor

As callbr is now allowed to have duplicate destinations, we can
have a callbr with a unique successor. Make sure it doesn't get
dropped, as we still need to preserve the side-effect.
---
 llvm/lib/Transforms/Utils/BasicBlockUtils.cpp      |  9 ++--
 .../Transforms/SimplifyCFG/callbr-destinations.ll  | 63 +++++++++++++++++-----
 2 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index 079b2fc..b08bf48 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -188,8 +188,10 @@ bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
 
   // Don't break self-loops.
   if (PredBB == BB) return false;
-  // Don't break unwinding instructions.
-  if (PredBB->getTerminator()->isExceptionalTerminator())
+
+  // Don't break unwinding instructions or terminators with other side-effects.
+  Instruction *PTI = PredBB->getTerminator();
+  if (PTI->isExceptionalTerminator() || PTI->mayHaveSideEffects())
     return false;
 
   // Can't merge if there are multiple distinct successors.
@@ -202,7 +204,7 @@ bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
   BasicBlock *NewSucc = nullptr;
   unsigned FallThruPath;
   if (PredecessorWithTwoSuccessors) {
-    if (!(PredBB_BI = dyn_cast<BranchInst>(PredBB->getTerminator())))
+    if (!(PredBB_BI = dyn_cast<BranchInst>(PTI)))
       return false;
     BranchInst *BB_JmpI = dyn_cast<BranchInst>(BB->getTerminator());
     if (!BB_JmpI || !BB_JmpI->isUnconditional())
@@ -256,7 +258,6 @@ bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
     Updates.push_back({DominatorTree::Delete, PredBB, BB});
   }
 
-  Instruction *PTI = PredBB->getTerminator();
   Instruction *STI = BB->getTerminator();
   Instruction *Start = &*BB->begin();
   // If there's nothing to move, mark the starting instruction as the last
diff --git a/llvm/test/Transforms/SimplifyCFG/callbr-destinations.ll b/llvm/test/Transforms/SimplifyCFG/callbr-destinations.ll
index aa1b37a..368bc5e 100644
--- a/llvm/test/Transforms/SimplifyCFG/callbr-destinations.ll
+++ b/llvm/test/Transforms/SimplifyCFG/callbr-destinations.ll
@@ -1,28 +1,67 @@
-; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -disable-output
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -S -simplifycfg -simplifycfg-require-and-preserve-domtree=1 | FileCheck %s
+
+define void @callbr_duplicate_dest() {
+; CHECK-LABEL: @callbr_duplicate_dest(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    callbr void asm sideeffect "", "!i"()
+; CHECK-NEXT:    to label [[BB:%.*]] [label %bb]
+; CHECK:       bb:
+; CHECK-NEXT:    ret void
 ;
-; Test that SimplifyCFG does not cause CallBr instructions to have duplicate
-; destinations, which will cause the verifier to assert.
+entry:
+  callbr void asm sideeffect "", "!i"()
+  to label %bb [label %bb]
+
+bb:
+  ret void
+}
 
-define void @fun0() {
+; TODO: Can fold to a duplicate callbr destination.
+define void @callbr_can_fold_to_duplicate_dest1() {
+; CHECK-LABEL: @callbr_can_fold_to_duplicate_dest1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    callbr void asm sideeffect "", "!i"()
+; CHECK-NEXT:    to label [[BB2:%.*]] [label %common.ret]
+; CHECK:       common.ret:
+; CHECK-NEXT:    ret void
+; CHECK:       bb2:
+; CHECK-NEXT:    br label [[COMMON_RET:%.*]]
+;
 entry:
   callbr void asm sideeffect "", "!i"()
-          to label %bb2 [label %bb1]
+  to label %bb2 [label %bb1]
 
-bb1:                                              ; preds = %bb
+bb1:
   ret void
 
-bb2:                                             ; preds = %bb
+bb2:
   ret void
 }
 
-define void @fun1() {
+; TODO: Can fold to a duplicate callbr destination.
+define void @callbr_can_fold_to_duplicate_dest2() {
+; CHECK-LABEL: @callbr_can_fold_to_duplicate_dest2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    callbr void asm sideeffect "", "!i,!i"()
+; CHECK-NEXT:    to label [[COMMON_RET:%.*]] [label [[BB2:%.*]], label %bb3]
+; CHECK:       common.ret:
+; CHECK-NEXT:    ret void
+; CHECK:       bb2:
+; CHECK-NEXT:    br label [[COMMON_RET]]
+; CHECK:       bb3:
+; CHECK-NEXT:    br label [[COMMON_RET]]
+;
 entry:
-  callbr void asm sideeffect "", "!i"()
-          to label %bb2 [label %bb1]
+  callbr void asm sideeffect "", "!i,!i"()
+  to label %bb1 [label %bb2, label %bb3]
+
+bb1:
+  ret void
 
-bb2:                                             ; preds = %bb
+bb2:
   ret void
 
-bb1:                                              ; preds = %bb
+bb3:
   ret void
 }
-- 
2.7.4