Fix SimplifyCFG::simplifyArgs for OSSA

Author: meg-gupta

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -3811,16 +3811,21 @@ bool SimplifyCFG::simplifyArgument(SILBasicBlock *BB, unsigned i) {
   // the uses in this block, and then rewrite the branch operands.
   LLVM_DEBUG(llvm::dbgs() << "unwrap argument:" << *A);
   A->replaceAllUsesWith(SILUndef::get(A->getType(), *BB->getParent()));
-  auto *NewArg =
-      BB->replacePhiArgument(i, proj->getType(), OwnershipKind::Owned);
+  auto *NewArg = BB->replacePhiArgument(i, proj->getType(),
+                                        BB->getArgument(i)->getOwnershipKind());
   proj->replaceAllUsesWith(NewArg);
 
   // Rewrite the branch operand for each incoming branch.
   for (auto *Pred : BB->getPredecessorBlocks()) {
     if (auto *Branch = cast<BranchInst>(Pred->getTerminator())) {
+      auto *BranchOpValue = cast<SingleValueInstruction>(Branch->getOperand(i));
       auto V = getInsertedValue(cast<SingleValueInstruction>(Branch->getArg(i)),
                                 proj);
       Branch->setOperand(i, V);
+      if (isInstructionTriviallyDead(BranchOpValue)) {
+        BranchOpValue->replaceAllUsesWithUndef();
+        BranchOpValue->eraseFromParent();
+      }
       addToWorklist(Pred);
     }
   }
@@ -3866,15 +3871,6 @@ bool SimplifyCFG::simplifyArgs(SILBasicBlock *BB) {
   if (BB->pred_empty())
     return false;
 
-  if (!EnableOSSARewriteTerminator && Fn.hasOwnership()) {
-    // TODO: OSSA phi support
-    if (llvm::any_of(BB->getArguments(), [this](SILArgument *arg) {
-          return !arg->getType().isTrivial(Fn);
-        })) {
-      return false;
-    }
-  }
-
   // Ignore blocks that are successors of terminators with mandatory args.
   for (SILBasicBlock *pred : BB->getPredecessorBlocks()) {
     if (hasMandatoryArgument(pred->getTerminator()))

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -270,6 +270,32 @@ struct ShrinkBorrowScopeTest : UnitTest {
   }
 };
 
+struct SimplifyCFGSimplifyArgument: UnitTest {
+  SimplifyCFGSimplifyArgument(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)
+        .simplifyArgument(arguments.takeBlock(), arguments.takeUInt());
+  }
+};
+
+struct SimplifyCFGSimplifyBlockArgs : UnitTest {
+  SimplifyCFGSimplifyBlockArgs(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)
+        .simplifyBlockArgs();
+  }
+};
+
 struct SimplifyCFGCanonicalizeSwitchEnum : UnitTest {
   SimplifyCFGCanonicalizeSwitchEnum(UnitTestRunner *pass) : UnitTest(pass) {}
   void invoke(Arguments &arguments) override {
@@ -444,6 +470,10 @@ void UnitTestRunner::withTest(StringRef name, Doit doit) {
     ADD_UNIT_TEST_SUBCLASS("shrink-borrow-scope", ShrinkBorrowScopeTest)
 
     // SimplifyCFG unit tests
+    ADD_UNIT_TEST_SUBCLASS("simplify-cfg-simplify-argument",
+                           SimplifyCFGSimplifyArgument)
+    ADD_UNIT_TEST_SUBCLASS("simplify-cfg-simplify-block-args",
+                           SimplifyCFGSimplifyBlockArgs)
     ADD_UNIT_TEST_SUBCLASS("simplify-cfg-canonicalize-switch-enum",
                            SimplifyCFGCanonicalizeSwitchEnum)
     ADD_UNIT_TEST_SUBCLASS("simplify-cfg-simplify-switch-enum-block",