Fix SimplifyCFG::simplifySwitchEnumOnObjcClassOptional for OSSA

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/Mandatory/DiagnoseInvalidEscapingCaptures.cpp

@@ -55,6 +55,18 @@ static bool checkNoEscapePartialApplyUse(Operand *oper, FollowUse followUses) {
       isa<CopyBlockWithoutEscapingInst>(user))
     return false;
 
+  // Look through copies, borrows, and conversions.
+  // getSingleValueCopyOrCast handles all result producing instructions for
+  // which onlyAffectsRefCount returns true.
+  if (SingleValueInstruction *copy = getSingleValueCopyOrCast(user)) {
+    // Only follow the copied operand. Other operands are incidental,
+    // as in the second operand of mark_dependence.
+    if (oper->getOperandNumber() == 0)
+      followUses(copy);
+
+    return false;
+  }
+
   // Ignore uses that are totally uninteresting. partial_apply [stack] is
   // terminated by a dealloc_stack instruction.
   if (isIncidentalUse(user) || onlyAffectsRefCount(user) ||
@@ -69,16 +81,6 @@ static bool checkNoEscapePartialApplyUse(Operand *oper, FollowUse followUses) {
       return false;
   }
 
-  // Look through copies, borrows, and conversions.
-  if (SingleValueInstruction *copy = getSingleValueCopyOrCast(user)) {
-    // Only follow the copied operand. Other operands are incidental,
-    // as in the second operand of mark_dependence.
-    if (oper->getOperandNumber() == 0)
-      followUses(copy);
-
-    return false;
-  }
-
   // Look through `differentiable_function`.
   if (auto *DFI = dyn_cast<DifferentiableFunctionInst>(user)) {
     followUses(DFI);

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,21 @@ 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);
+    auto *someBBArg = someBB->getArgument(0);
+    if (!someBBArg->use_empty()) {
+      assert(optionalPayload != someBBArg);
+      someBBArg->replaceAllUsesWith(payloadCast);
+    }
+    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) {}
@@ -515,6 +528,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)

test/SILOptimizer/outliner.swift

@@ -63,7 +63,6 @@ public func testOutlining() {
 // CHECK:  [[OBJ:%.*]] = open_existential_ref {{.*}} : $AnyObject to $@opened("{{.*}}", AnyObject) Self
 // CHECK:  [[METH:%.*]] = objc_method [[OBJ]] : $@opened("{{.*}}", AnyObject) Self, #Treeish.treeishChildren!foreign : <Self where Self : Treeish> (Self) -> () -> [Any]?
 // CHECK:  [[RES:%.*]] = apply [[METH]]([[OBJ]]) : $@convention(objc_method)
-// CHECK:  switch_enum [[RES]]
 // CHECK: } // end sil function '$s8outliner9dontCrash1ayyp_tF'
 
 // CHECK-LABEL: sil shared [noinline] @$sSo5GizmoC14stringPropertySSSgvgToTeab_ : {{.*}} {