[SimplifyCFG] Consider preds to switch in simplifyDuplicateSwitchArms

Allow a duplicate basic block with multiple predecessors to the
jump table to be simplified, by considering that the same basic
block may appear in more switch cases.

Author: antoniofrighetto

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -7474,9 +7474,6 @@ static bool simplifySwitchOfCmpIntrinsic(SwitchInst *SI, IRBuilderBase &Builder,
 /// IncomingValue and add it in the Wrapper so isEqual can do O(1) checking
 /// of the incoming values.
 struct SwitchSuccWrapper {
-  // Keep so we can use SwitchInst::setSuccessor to do the replacement. It won't
-  // be important to equality though.
-  unsigned SuccNum;
   BasicBlock *Dest;
   DenseMap<PHINode *, SmallDenseMap<BasicBlock *, Value *, 8>> *PhiPredIVs;
 };
@@ -7563,6 +7560,7 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
   SmallPtrSet<PHINode *, 8> Phis;
   SmallPtrSet<BasicBlock *, 8> Seen;
   DenseMap<PHINode *, SmallDenseMap<BasicBlock *, Value *, 8>> PhiPredIVs;
+  DenseMap<BasicBlock *, SmallVector<unsigned, 4>> BBToSuccessorIndexes;
   SmallVector<SwitchSuccWrapper> Cases;
   Cases.reserve(SI->getNumSuccessors());
 
@@ -7575,8 +7573,9 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
       continue;
 
     // FIXME: This case needs some extra care because the terminators other than
-    // SI need to be updated.
-    if (BB->hasNPredecessorsOrMore(2))
+    // SI need to be updated. For now, consider only backedges to the SI.
+    if (BB->hasNPredecessorsOrMore(4) ||
+        BB->getUniquePredecessor() != SI->getParent())
       continue;
 
     // FIXME: Relax that the terminator is a BranchInst by checking for equality
@@ -7591,8 +7590,11 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
       for (BasicBlock *Succ : BI->successors())
         for (PHINode &Phi : Succ->phis())
           Phis.insert(&Phi);
+      // Add the successor only if not previously visited.
+      Cases.emplace_back(SwitchSuccWrapper{BB, &PhiPredIVs});
     }
-    Cases.emplace_back(SwitchSuccWrapper{I, BB, &PhiPredIVs});
+
+    BBToSuccessorIndexes[BB].emplace_back(I);
   }
 
   // Precompute a data structure to improve performance of isEqual for
@@ -7627,7 +7629,9 @@ bool SimplifyCFGOpt::simplifyDuplicateSwitchArms(SwitchInst *SI,
       // We know that SI's parent BB no longer dominates the old case successor
       // since we are making it dead.
       Updates.push_back({DominatorTree::Delete, SI->getParent(), SSW.Dest});
-      SI->setSuccessor(SSW.SuccNum, (*It)->Dest);
+      const auto &Successors = BBToSuccessorIndexes.at(SSW.Dest);
+      for (unsigned Idx : Successors)
+        SI->setSuccessor(Idx, (*It)->Dest);
       MadeChange = true;
     }
   }

llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll

@@ -320,14 +320,14 @@ define i32 @overflow(i32 %type) {
 ; CHECK-NEXT:      i32 1, label [[IF_END:%.*]]
 ; CHECK-NEXT:      i32 2, label [[SW_BB2:%.*]]
 ; CHECK-NEXT:    ]
-; CHECK:       sw.bb1:
-; CHECK-NEXT:    br label [[SW_DEFAULT]]
 ; CHECK:       sw.bb2:
-; CHECK-NEXT:    br label [[SW_DEFAULT]]
+; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       sw.bb3:
-; CHECK-NEXT:    br label [[SW_DEFAULT]]
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
-; CHECK-NEXT:    [[DIRENT_TYPE_0:%.*]] = phi i32 [ 6, [[SW_BB3]] ], [ 5, [[SW_BB2]] ], [ 0, [[IF_END]] ], [ 3, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[DIRENT_TYPE_0:%.*]] = phi i32 [ 3, [[SW_DEFAULT]] ], [ 6, [[SW_BB3]] ], [ 5, [[SW_BB2]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    ret i32 [[DIRENT_TYPE_0]]
 ;
 entry:

llvm/test/Transforms/SimplifyCFG/switch-dup-bbs.ll

@@ -136,17 +136,13 @@ define i32 @switch_dup_exit(i32 %val) {
 ; SIMPLIFY-CFG-NEXT:      i32 1, label %[[EXIT:.*]]
 ; SIMPLIFY-CFG-NEXT:      i32 11, label %[[EXIT]]
 ; SIMPLIFY-CFG-NEXT:      i32 22, label %[[BB1:.*]]
-; SIMPLIFY-CFG-NEXT:      i32 15, label %[[BB2:.*]]
-; SIMPLIFY-CFG-NEXT:      i32 0, label %[[BB2]]
 ; SIMPLIFY-CFG-NEXT:    ]
 ; SIMPLIFY-CFG:       [[BB1]]:
 ; SIMPLIFY-CFG-NEXT:    br label %[[EXIT]]
-; SIMPLIFY-CFG:       [[BB2]]:
-; SIMPLIFY-CFG-NEXT:    br label %[[EXIT]]
 ; SIMPLIFY-CFG:       [[DEFAULT]]:
 ; SIMPLIFY-CFG-NEXT:    br label %[[EXIT]]
 ; SIMPLIFY-CFG:       [[EXIT]]:
-; SIMPLIFY-CFG-NEXT:    [[RET:%.*]] = phi i32 [ 0, %[[DEFAULT]] ], [ 0, %[[BB2]] ], [ 3, %[[BB1]] ], [ 1, %[[ENTRY]] ], [ 1, %[[ENTRY]] ]
+; SIMPLIFY-CFG-NEXT:    [[RET:%.*]] = phi i32 [ 0, %[[DEFAULT]] ], [ 3, %[[BB1]] ], [ 1, %[[ENTRY]] ], [ 1, %[[ENTRY]] ]
 ; SIMPLIFY-CFG-NEXT:    ret i32 [[RET]]
 ;
 entry:
@@ -175,19 +171,11 @@ exit:
 define i64 @switch_dup_exit_2(i32 %val) {
 ; SIMPLIFY-CFG-LABEL: define i64 @switch_dup_exit_2(
 ; SIMPLIFY-CFG-SAME: i32 [[VAL:%.*]]) {
-; SIMPLIFY-CFG-NEXT:  [[ENTRY:.*]]:
-; SIMPLIFY-CFG-NEXT:    switch i32 [[VAL]], label %[[DEFAULT:.*]] [
-; SIMPLIFY-CFG-NEXT:      i32 1, label %[[EXIT:.*]]
-; SIMPLIFY-CFG-NEXT:      i32 11, label %[[EXIT]]
-; SIMPLIFY-CFG-NEXT:      i32 13, label %[[BB1:.*]]
-; SIMPLIFY-CFG-NEXT:      i32 0, label %[[BB1]]
-; SIMPLIFY-CFG-NEXT:    ]
-; SIMPLIFY-CFG:       [[BB1]]:
-; SIMPLIFY-CFG-NEXT:    br label %[[EXIT]]
-; SIMPLIFY-CFG:       [[DEFAULT]]:
-; SIMPLIFY-CFG-NEXT:    br label %[[EXIT]]
-; SIMPLIFY-CFG:       [[EXIT]]:
-; SIMPLIFY-CFG-NEXT:    [[RET:%.*]] = phi i64 [ 0, %[[DEFAULT]] ], [ 0, %[[BB1]] ], [ 1, %[[ENTRY]] ], [ 1, %[[ENTRY]] ]
+; SIMPLIFY-CFG-NEXT:  [[ENTRY:.*:]]
+; SIMPLIFY-CFG-NEXT:    [[SWITCH_SELECTCMP_CASE1:%.*]] = icmp eq i32 [[VAL]], 1
+; SIMPLIFY-CFG-NEXT:    [[SWITCH_SELECTCMP_CASE2:%.*]] = icmp eq i32 [[VAL]], 11
+; SIMPLIFY-CFG-NEXT:    [[SWITCH_SELECTCMP:%.*]] = or i1 [[SWITCH_SELECTCMP_CASE1]], [[SWITCH_SELECTCMP_CASE2]]
+; SIMPLIFY-CFG-NEXT:    [[RET:%.*]] = select i1 [[SWITCH_SELECTCMP]], i64 1, i64 0
 ; SIMPLIFY-CFG-NEXT:    ret i64 [[RET]]
 ;
 entry: