[SwitchLowering] Handle multiple destinations on condensed case stmts

[bruno]

Summary

Switch cases statements with sequential values that branch to the same
destination BB may often be handled together in a single new source
BB. In this scenario we need to remove remaining incoming values from
PHI instructions in the destination BB, as to match the number of source
branches.

[ributzka]

Hi Bruno,

thanks for tracking this one down.

LGTM

--Juergen

[bruno]

Committed in r222926

lib/Transforms/Utils/LowerSwitch.cpp

for (LowerSwitch::CaseVector::const_iterator B = C.begin(),
          E = C.end(); B != E; ) {
     O << *B->Low << " -" << *B->High;
     if (++B != E) O << ", ";
   }
 
   return O << "]";
 }
 
-/// \brief Update the first occurrence of the "switch statement" BB in the PHI
-/// node with the "new" BB. The other occurrences will be updated by subsequent
-/// calls to this function.
-///
-/// Switch statements may have more than one incoming edge into the same BB if
-/// they all have the same value. When the switch statement is converted these
-/// incoming edges are now coming from multiple BBs.
-static void fixPhis(BasicBlock *SuccBB, BasicBlock *OrigBB, BasicBlock *NewBB) {
-  for (BasicBlock::iterator I = SuccBB->begin(), E = SuccBB->getFirstNonPHI();
-       I != E; ++I) {
+// \brief Update the first occurrence of the "switch statement" BB in the PHI
+// node with the "new" BB. The other occurrences will:
+//
+// 1) Be updated by subsequent calls to this function.  Switch statements may
+// have more than one outcoming edge into the same BB if they all have the same
+// value. When the switch statement is converted these incoming edges are now
+// coming from multiple BBs.
+// 2) Removed if subsequent incoming values now share the same case, i.e.,
+// multiple outcome edges are condensed into one. This is necessary to keep the
+// number of phi values equal to the number of branches to SuccBB.
+static void fixPhis(BasicBlock *SuccBB, BasicBlock *OrigBB, BasicBlock *NewBB,
+                    unsigned NumMergedCases) {
+  for (BasicBlock::iterator I = SuccBB->begin(), IE = SuccBB->getFirstNonPHI();
+       I != IE; ++I) {
     PHINode *PN = cast<PHINode>(I);
 
     // Only update the first occurence.
-    for (unsigned Idx = 0, E = PN->getNumIncomingValues(); Idx != E; ++Idx) {
+    unsigned Idx = 0, E = PN->getNumIncomingValues();
+    for (; Idx != E; ++Idx) {
       if (PN->getIncomingBlock(Idx) == OrigBB) {
         PN->setIncomingBlock(Idx, NewBB);
         break;
       }
     }
+
+    // Remove additional occurences coming from condensed cases and keep the
+    // number of incoming values equal to the number of branches to SuccBB.
+    for (++Idx; NumMergedCases > 0 && Idx != E; ++Idx)
+      if (PN->getIncomingBlock(Idx) == OrigBB) {
+        PN->removeIncomingValue(Idx);
+        NumMergedCases--;
+      }
   }
 }
 
 // switchConvert - Convert the switch statement into a binary lookup of
 // the case values. The function recursively builds this tree.
 // LowerBound and UpperBound are used to keep track of the bounds for Val
 // that have already been checked by a block emitted by one of the previous
 // calls to switchConvert in the call stack.
 BasicBlock *LowerSwitch::switchConvert(CaseItr Begin, CaseItr End,
                                        ConstantInt *LowerBound,
                                        ConstantInt *UpperBound, Value *Val,
                                        BasicBlock *Predecessor,
                                        BasicBlock *OrigBlock,
                                        BasicBlock *Default) {
   unsigned Size = End - Begin;
 
   if (Size == 1) {
     // Check if the Case Range is perfectly squeezed in between
     // already checked Upper and Lower bounds. If it is then we can avoid
     // emitting the code that checks if the value actually falls in the range
     // because the bounds already tell us so.
     if (Begin->Low == LowerBound && Begin->High == UpperBound) {
-      fixPhis(Begin->BB, OrigBlock, Predecessor);
+      unsigned NumMergedCases = 0;
+      if (LowerBound && UpperBound)
+        NumMergedCases =
+            UpperBound->getSExtValue() - LowerBound->getSExtValue();
+      fixPhis(Begin->BB, OrigBlock, Predecessor, NumMergedCases);
       return Begin->BB;
     }
     return newLeafBlock(*Begin, Val, OrigBlock, Default);
   }
 
   unsigned Mid = Size / 2;
   std::vector<CaseRange> LHS(Begin, Begin + Mid);
   DEBUG(dbgs() << "LHS: " << LHS << "\n");

test/Transforms/Util/lowerswitch.ll

 ; RUN: opt -lowerswitch -S < %s | FileCheck %s
 
 ; Test that we don't crash and have a different basic block for each incoming edge.
-define void @test_lower_switch() {
-; CHECK-LABEL: @test_lower_switch
+define void @test0() {
+; CHECK-LABEL: @test0
 ; CHECK: %merge = phi i64 [ 1, %BB3 ], [ 0, %NewDefault ], [ 0, %NodeBlock5 ], [ 0, %LeafBlock1 ]
 BB1:
   switch i32 undef, label %BB2 [
     i32 3, label %BB2
     i32 5, label %BB2
     i32 0, label %BB3
     i32 2, label %BB3
     i32 4, label %BB3
   ]
 
 BB2:
   %merge = phi i64 [ 1, %BB3 ], [ 0, %BB1 ], [ 0, %BB1 ], [ 0, %BB1 ]
   ret void
 
 BB3:
   br label %BB2
 }
+
+; Test switch cases that are merged into a single case during lowerswitch
+; (take 84 and 85 below) - check that the number of incoming phi values match
+; the number of branches.
+define void @test1() {
+; CHECK-LABEL: @test1
+entry:
+  br label %bb1
+
+bb1:
+  switch i32 undef, label %bb1 [
+    i32 84, label %bb3
+    i32 85, label %bb3
+    i32 86, label %bb2
+    i32 78, label %exit
+    i32 99, label %bb3
+  ]
+
+bb2:
+  br label %bb3
+
+bb3:
+; CHECK-LABEL: bb3
+; CHECL-NEXT: %tmp = phi i32 [ 1, %NodeBlock ], [ 0, %bb2 ], [ 1, %LeafBlock3 ]
+  %tmp = phi i32 [ 1, %bb1 ], [ 0, %bb2 ], [ 1, %bb1 ], [ 1, %bb1 ]
+  br label %exit
+
+exit:
+  ret void
+}