Merge pull request #4153 from swiftix/swift-3.0-branch

[sil-combine] Fix a bug in the cond_br (select_enum) -> switch_enum peephole

Author: ematejska

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -997,35 +997,80 @@ SILInstruction *SILCombiner::visitCondBranchInst(CondBranchInst *CBI) {
     if (!CBI->getFalseBB()->getSinglePredecessor())
       return nullptr;
 
-    SILBasicBlock *Default = nullptr;
-
+    SILBasicBlock *DefaultBB = nullptr;
     match_integer<0> Zero;
 
     if (SEI->hasDefault()) {
+      // Default result should be an integer constant.
+      if (!isa<IntegerLiteralInst>(SEI->getDefaultResult()))
+        return nullptr;
       bool isFalse = match(SEI->getDefaultResult(), Zero);
-      Default = isFalse ? CBI->getFalseBB() : Default = CBI->getTrueBB();
+      // Pick the default BB.
+      DefaultBB = isFalse ? CBI->getFalseBB() : CBI->getTrueBB();
+    }
+
+    if (!DefaultBB) {
+      // Find the targets for the majority of cases and pick it
+      // as a default BB.
+      unsigned TrueBBCases = 0;
+      unsigned FalseBBCases = 0;
+      for (int i = 0, e = SEI->getNumCases(); i < e; ++i) {
+        auto Pair = SEI->getCase(i);
+        if (isa<IntegerLiteralInst>(Pair.second)) {
+          bool isFalse = match(Pair.second, Zero);
+          if (!isFalse) {
+            TrueBBCases++;
+          } else {
+            FalseBBCases++;
+          }
+          continue;
+        }
+        return nullptr;
+      }
+
+      if (FalseBBCases > TrueBBCases)
+        DefaultBB = CBI->getFalseBB();
+      else
+        DefaultBB = CBI->getTrueBB();
     }
 
-    // We can now convert cond_br(select_enum) into switch_enum
+    assert(DefaultBB && "Default should be defined at this point");
+
+    unsigned NumTrueBBCases = 0;
+    unsigned NumFalseBBCases = 0;
+
+    if (DefaultBB == CBI->getFalseBB())
+      NumFalseBBCases++;
+    else
+      NumTrueBBCases++;
+
+    // We can now convert cond_br(select_enum) into switch_enum.
     SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 8> Cases;
     for (int i = 0, e = SEI->getNumCases(); i < e; ++i) {
       auto Pair = SEI->getCase(i);
-      if (isa<IntegerLiteralInst>(Pair.second)) {
-        bool isFalse = match(Pair.second, Zero);
-        if (!isFalse && Default != CBI->getTrueBB()) {
-          Cases.push_back(std::make_pair(Pair.first, CBI->getTrueBB()));
-        }
-        if (isFalse && Default != CBI->getFalseBB()) {
-          Cases.push_back(std::make_pair(Pair.first, CBI->getFalseBB()));
-        }
-        continue;
+
+      // Bail if one of the results is not an integer constant.
+      if (!isa<IntegerLiteralInst>(Pair.second))
+        return nullptr;
+
+      // Add a switch case.
+      bool isFalse = match(Pair.second, Zero);
+      if (!isFalse && DefaultBB != CBI->getTrueBB()) {
+        Cases.push_back(std::make_pair(Pair.first, CBI->getTrueBB()));
+        NumTrueBBCases++;
+      }
+      if (isFalse && DefaultBB != CBI->getFalseBB()) {
+        Cases.push_back(std::make_pair(Pair.first, CBI->getFalseBB()));
+        NumFalseBBCases++;
       }
+    }
 
+    // Bail if a switch_enum would introduce a critical edge.
+    if (NumTrueBBCases > 1 || NumFalseBBCases > 1)
       return nullptr;
-    }
 
     return Builder.createSwitchEnum(SEI->getLoc(), SEI->getEnumOperand(),
-                                    Default, Cases);
+                                    DefaultBB, Cases);
   }
 
   return nullptr;

test/SILOptimizer/sil_combine_enums.sil

@@ -9,6 +9,13 @@ class SomeClass {
   func hash() -> Int
 }
 
+enum Numerals {
+  case One
+  case Two
+  case Three
+  case Four
+}
+
 sil @external_func: $@convention(thin) () -> ()
 
 //CHECK-LABEL: eliminate_sw_enum_addr
@@ -177,6 +184,68 @@ bb2:
   unreachable
 }
 
+// Check that cond_br(select_enum) is converted into switch_enum.
+// CHECK-LABEL: sil @convert_select_enum_cond_br_to_switch_enum2
+// CHECK: bb0
+// CHECK-NOT: select_enum
+// CHECK-NOT: return
+// CHECK: switch_enum %0 : $Numerals, case #Numerals.Two!enumelt: bb3, default bb2
+// CHECK: return
+// CHECK: }
+sil @convert_select_enum_cond_br_to_switch_enum2 : $@convention(thin) (Numerals) -> Builtin.Int64 {
+bb0(%0 : $Numerals):
+  %2 = integer_literal $Builtin.Int1, 0
+  %3 = integer_literal $Builtin.Int1, -1
+  // All cases but one are the same. So, they can be made a default for the switch_enum.
+  %4 = select_enum %0 : $Numerals, case #Numerals.One!enumelt: %3, case #Numerals.Two!enumelt: %2, case #Numerals.Three!enumelt: %3, case #Numerals.Four!enumelt: %3 : $Builtin.Int1
+
+  cond_br %4, bb2, bb3
+
+bb1:
+  %7 = integer_literal $Builtin.Int64, 10
+  return %7 : $Builtin.Int64
+
+bb2:
+  %10 = function_ref @external_func: $@convention(thin) () -> ()
+  apply %10(): $@convention(thin) () -> ()
+  br bb1
+
+bb3:
+  br bb1
+}
+
+// Check that cond_br(select_enum) is converted into switch_enum.
+// This test checks that select_enum instructions with default cases are handled correctly.
+// CHECK-LABEL: sil @convert_select_enum_cond_br_to_switch_enum3
+// CHECK: bb0
+// CHECK-NOT: select_enum
+// CHECK-NOT: return
+// CHECK: switch_enum %0 : $Numerals, case #Numerals.Two!enumelt: bb3, default bb2
+// CHECK: return
+// CHECK: }
+sil @convert_select_enum_cond_br_to_switch_enum3 : $@convention(thin) (Numerals) -> Builtin.Int64 {
+bb0(%0 : $Numerals):
+  %2 = integer_literal $Builtin.Int1, 0
+  %3 = integer_literal $Builtin.Int1, -1
+  // There is only one case, whose result is different from default and other cases
+  // Thus all other cases can be folded into a default cases of a switch_enum.
+  %4 = select_enum %0 : $Numerals, case #Numerals.One!enumelt: %3, case #Numerals.Two!enumelt: %2, case #Numerals.Three!enumelt: %3, default %3 : $Builtin.Int1
+
+  cond_br %4, bb2, bb3
+
+bb1:
+  %7 = integer_literal $Builtin.Int64, 10
+  return %7 : $Builtin.Int64
+
+bb2:
+  %10 = function_ref @external_func: $@convention(thin) () -> ()
+  apply %10(): $@convention(thin) () -> ()
+  br bb1
+
+bb3:
+  br bb1
+}
+
 
 // Check that cond_br(select_enum) is not converted into switch_enum as it would create a critical edge, which
 // is not originating from cond_br/br. And this is forbidden in a canonical SIL form.
@@ -206,6 +275,64 @@ bb2:
   br bb1
 }
 
+// Check that cond_br(select_enum) is not converted into switch_enum as it would create a critical edge, which
+// is not originating from cond_br/br. And this is forbidden in a canonical SIL form.
+//
+// CHECK-LABEL: sil @dont_convert_select_enum_cond_br_to_switch_enum2
+// CHECK: select_enum
+// CHECK-NOT: switch_enum
+// CHECK: return
+sil @dont_convert_select_enum_cond_br_to_switch_enum2 : $@convention(thin) (Numerals) -> Builtin.Int64 {
+bb0(%0 : $Numerals):
+  %2 = integer_literal $Builtin.Int1, 0
+  %3 = integer_literal $Builtin.Int1, -1
+  // There are two cases for each possible outcome.
+  // This means that we would always get a critical edge, if we convert it into a switch_enum.
+  %4 = select_enum %0 : $Numerals, case #Numerals.One!enumelt: %3, case #Numerals.Two!enumelt: %2, case #Numerals.Three!enumelt: %3, case #Numerals.Four!enumelt: %2 : $Builtin.Int1
+
+  cond_br %4, bb2, bb3
+
+bb1:
+  %7 = integer_literal $Builtin.Int64, 10
+  return %7 : $Builtin.Int64
+
+bb2:
+  %10 = function_ref @external_func: $@convention(thin) () -> ()
+  apply %10(): $@convention(thin) () -> ()
+  br bb1
+
+bb3:
+  br bb1
+}
+
+// Check that cond_br(select_enum) is not converted into switch_enum,
+// because the result of the default case is not an integer literal.
+// CHECK-LABEL: sil @dont_convert_select_enum_cond_br_to_switch_enum3
+// CHECK: select_enum
+// CHECK-NOT: switch_enum
+// CHECK: return
+sil @dont_convert_select_enum_cond_br_to_switch_enum3 : $@convention(thin) (Numerals, Builtin.Int1) -> Builtin.Int64 {
+bb0(%0 : $Numerals, %1 : $Builtin.Int1):
+  %2 = integer_literal $Builtin.Int1, 0
+  %3 = integer_literal $Builtin.Int1, -1
+  // All cases but one are the same. So, they can be made a default for the switch_enum.
+  %4 = select_enum %0 : $Numerals, case #Numerals.One!enumelt: %3, case #Numerals.Two!enumelt: %2, case #Numerals.Three!enumelt: %3, default %1 : $Builtin.Int1
+
+  cond_br %4, bb2, bb3
+
+bb1:
+  %7 = integer_literal $Builtin.Int64, 10
+  return %7 : $Builtin.Int64
+
+bb2:
+  %10 = function_ref @external_func: $@convention(thin) () -> ()
+  apply %10(): $@convention(thin) () -> ()
+  br bb1
+
+bb3:
+  br bb1
+}
+
 public class C {}
 public struct S {}
 public struct T {