Enable SimplifyCFG::simplifySwitchEnumOnObjcClassOptional

Author: meg-gupta

lib/SIL/Utils/InstructionUtils.cpp

@@ -279,8 +279,9 @@ bool swift::onlyAffectsRefCount(SILInstruction *user) {
   switch (user->getKind()) {
   default:
     return false;
-  case SILInstructionKind::AutoreleaseValueInst:
+  case SILInstructionKind::CopyValueInst:
   case SILInstructionKind::DestroyValueInst:
+  case SILInstructionKind::AutoreleaseValueInst:
   case SILInstructionKind::ReleaseValueInst:
   case SILInstructionKind::RetainValueInst:
   case SILInstructionKind::StrongReleaseInst:

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -1847,7 +1847,9 @@ static bool containsOnlyObjMethodCallOnOptional(SILValue optionalValue,
 
     // The branch should forward one of the objc_method call.
     if (auto *br = dyn_cast<BranchInst>(inst)) {
-      if (br->getNumArgs() == 0 || br->getNumArgs() > 1)
+      if (br->getNumArgs() == 0)
+        continue;
+      if (br->getNumArgs() > 1)
         return false;
       auto branchArg = br->getArg(0);
       if (std::find(objCApplies.begin(), objCApplies.end(), branchArg) ==
@@ -1895,6 +1897,9 @@ static bool onlyForwardsNone(SILBasicBlock *noneBB, SILBasicBlock *someBB,
       continue;
     }
     if (auto *noneBranch = dyn_cast<BranchInst>(inst)) {
+      if (noneBranch->getNumArgs() == 0) {
+        continue;
+      }
       if (noneBranch->getNumArgs() != 1 ||
           (noneBranch->getArg(0) != SEI->getOperand() &&
            noneBranch->getArg(0) != optionalNone))
@@ -2006,12 +2011,6 @@ static bool hasSameUltimateSuccessor(SILBasicBlock *noneBB, SILBasicBlock *someB
 ///    %4 = enum #Optional.none
 ///    br mergeBB(%4)
 bool SimplifyCFG::simplifySwitchEnumOnObjcClassOptional(SwitchEnumInst *SEI) {
-  // TODO: OSSA; handle non-trivial enum case cleanup
-  // (simplify_switch_enum_objc.sil).
-  if (!EnableOSSARewriteTerminator && Fn.hasOwnership()) {
-    return false;
-  }
-
   auto optional = SEI->getOperand();
   auto optionalPayloadType = optional->getType().getOptionalObjectType();
   if (!optionalPayloadType ||
@@ -2048,10 +2047,16 @@ bool SimplifyCFG::simplifySwitchEnumOnObjcClassOptional(SwitchEnumInst *SEI) {
                                                      optionalPayloadType);
   optionalPayload->replaceAllUsesWith(payloadCast);
   auto *switchBB = SEI->getParent();
-  if (someBB->getNumArguments())
-    Builder.createBranch(SEI->getLoc(), someBB, SILValue(payloadCast));
-  else
+
+  if (!someBB->args_empty()) {
+    assert(someBB->getNumArguments() == 1);
+    someBB->eraseArgument(0);
     Builder.createBranch(SEI->getLoc(), someBB);
+  }
+  else {
+    assert(!Fn.hasOwnership());
+    Builder.createBranch(SEI->getLoc(), someBB);
+  }
 
   SEI->eraseFromParent();
   addToWorklist(switchBB);

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -321,6 +321,19 @@ struct SimplifyCFGSimplifySwitchEnumBlock : UnitTest {
   }
 };
 
+struct SimplifyCFGSwitchEnumOnObjcClassOptional : UnitTest {
+  SimplifyCFGSwitchEnumOnObjcClassOptional(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)
+        .simplifySwitchEnumOnObjcClassOptional(
+            cast<SwitchEnumInst>(arguments.takeInstruction()));
+  }
+};
+
 struct SimplifyCFGSimplifySwitchEnumUnreachableBlocks : UnitTest {
   SimplifyCFGSimplifySwitchEnumUnreachableBlocks(UnitTestRunner *pass)
       : UnitTest(pass) {}
@@ -481,6 +494,9 @@ void UnitTestRunner::withTest(StringRef name, Doit doit) {
     ADD_UNIT_TEST_SUBCLASS(
         "simplify-cfg-simplify-switch-enum-unreachable-blocks",
         SimplifyCFGSimplifySwitchEnumUnreachableBlocks)
+    ADD_UNIT_TEST_SUBCLASS(
+        "simplify-cfg-simplify-switch-enum-on-objc-class-optional",
+        SimplifyCFGSwitchEnumOnObjcClassOptional)
     ADD_UNIT_TEST_SUBCLASS(
         "simplify-cfg-simplify-term-with-identical-dest-blocks",
         SimplifyCFGSimplifyTermWithIdenticalDestBlocks)