Enable SimplifyCFG::canonicalizeSwitchEnums and add unit tests

Author: meg-gupta

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -2748,9 +2748,6 @@ bool SimplifyCFG::simplifyBlocks() {
 /// Canonicalize all switch_enum and switch_enum_addr instructions.
 /// If possible, replace the default with the corresponding unique case.
 bool SimplifyCFG::canonicalizeSwitchEnums() {
-  if (!EnableOSSASimplifyCFG && Fn.hasOwnership()) {
-    return false;
-  }
   bool Changed = false;
   for (auto &BB : Fn) {
     TermInst *TI = BB.getTerminator();
@@ -2768,13 +2765,7 @@ bool SimplifyCFG::canonicalizeSwitchEnums() {
     NullablePtr<EnumElementDecl> defaultDecl = SWI.getUniqueCaseForDefault();
     if (!defaultDecl)
       continue;
-    
-    if (!EnableOSSARewriteTerminator && Fn.hasOwnership()) {
-      if (!SWI.getOperand()->getType().isTrivial(Fn)) {
-        // TODO: Test and enable this case.
-        continue;
-      }
-    }
+
     LLVM_DEBUG(llvm::dbgs() << "simplify canonical switch_enum\n");
 
     // Construct a new instruction by copying all the case entries.

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -75,6 +75,7 @@
 #include "swift/SILOptimizer/Analysis/BasicCalleeAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Transforms/SimplifyCFG.h"
 #include "swift/SILOptimizer/Utils/CanonicalizeBorrowScope.h"
 #include "swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -303,6 +304,18 @@ struct DumpFunction : UnitTest {
   void invoke(Arguments &arguments) override { getFunction()->dump(); }
 };
 
+struct SimplifyCFGCanonicalizeSwitchEnum : UnitTest {
+  SimplifyCFGCanonicalizeSwitchEnum(UnitTestRunner *pass) : UnitTest(pass) {}
+  void invoke(Arguments &arguments) override {
+    auto *passToRun = cast<SILFunctionTransform>(createSimplifyCFG());
+    passToRun->injectPassManager(getPass()->getPassManager());
+    passToRun->injectFunction(getFunction());
+    SimplifyCFG(*getFunction(), *passToRun, /*VerifyAll=*/false,
+                /*EnableJumpThread=*/false)
+        .canonicalizeSwitchEnums();
+  }
+};
+
 // Arguments: NONE
 // Dumps: the index of the self argument of the current function
 struct FunctionGetSelfArgumentIndex : UnitTest {
@@ -386,6 +399,8 @@ void UnitTestRunner::withTest(StringRef name, Doit doit) {
         PrunedLivenessBoundaryWithListOfLastUsersInsertionPointsTest)
     ADD_UNIT_TEST_SUBCLASS("shrink-borrow-scope", ShrinkBorrowScopeTest)
     ADD_UNIT_TEST_SUBCLASS("is-deinit-barrier", IsDeinitBarrierTest)
+    ADD_UNIT_TEST_SUBCLASS("simplify-cfg-canonicalize-switch-enum",
+                           SimplifyCFGCanonicalizeSwitchEnum)
     /// [new_tests] Add the new mapping from string to subclass above this line.
 
 #undef ADD_UNIT_TEST_SUBCLASS

test/SILOptimizer/simplify_cfg_ossa_switch_enum.sil

@@ -0,0 +1,84 @@
+// RUN: %target-sil-opt -unit-test-runner %s 2>&1 | %FileCheck %s
+
+class Klass {
+}
+
+enum FakeOptional<T> {
+  case some(T)
+  case none
+}
+
+enum E1 {
+  case A(Klass)
+  case B(Klass)
+  case C(Klass)
+}
+
+enum E2 {
+  case A(Klass)
+  case B(Klass)
+  case C
+}
+
+// CHECK-LABEL: sil [ossa] @test_canonicalize_switch_enum1 :
+// CHECK-NOT: default
+// CHECK-LABEL: } // end sil function 'test_canonicalize_switch_enum1'
+sil [ossa] @test_canonicalize_switch_enum1 : $@convention(thin) (@owned FakeOptional<Klass>) -> () {
+bb0(%0 : @owned $FakeOptional<Klass>):
+  test_specification "simplify-cfg-canonicalize-switch-enum"
+  switch_enum %0 : $FakeOptional<Klass>, case #FakeOptional.none!enumelt: bb1, default bb2
+
+bb1:
+  br bb3
+
+bb2(%5 : @owned $Klass):
+  destroy_value %5 : $Klass
+  br bb3
+
+bb3:
+  %t = tuple ()
+  return %t : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_canonicalize_switch_enum2 :
+// CHECK: default
+// CHECK-LABEL: } // end sil function 'test_canonicalize_switch_enum2'
+sil [ossa] @test_canonicalize_switch_enum2 : $@convention(thin) (@owned E1) -> () {
+bb0(%0 : @owned $E1):
+  test_specification "simplify-cfg-canonicalize-switch-enum"
+  switch_enum %0 : $E1, case #E1.A!enumelt: bb1, default bb2
+
+bb1(%2 : @owned $Klass):
+  destroy_value %2 : $Klass
+  br bb3
+
+bb2(%5 : @owned $E1):
+  destroy_value %5 : $E1
+  br bb3
+
+bb3:
+  %t = tuple ()
+  return %t : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_canonicalize_switch_enum3 :
+// CHECK: default
+// CHECK-LABEL: } // end sil function 'test_canonicalize_switch_enum3'
+sil [ossa] @test_canonicalize_switch_enum3 : $@convention(thin) (@owned E2) -> () {
+bb0(%0 : @owned $E2):
+  test_specification "simplify-cfg-canonicalize-switch-enum"
+  switch_enum %0 : $E2, case #E2.A!enumelt: bb1, default bb2
+
+bb1(%2 : @owned $Klass):
+  destroy_value %2 : $Klass
+  br bb3
+
+bb2(%5 : @owned $E2):
+  destroy_value %5 : $E2
+  br bb3
+
+bb3:
+  %t = tuple ()
+  return %t : $()
+}
+