Fix simplifyArguments for a few edges cases in OSSA

Author: meg-gupta

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -233,7 +233,8 @@ bool SimplifyCFG::threadEdge(const ThreadInfo &ti) {
         Builder.createBranch(SEI->getLoc(), ThreadedSuccessorBlock, {UED});
       } else {
         assert(SEI->getDefaultBB() == ThreadedSuccessorBlock);
-        Builder.createBranch(SEI->getLoc(), ThreadedSuccessorBlock, SEI->getOperand());
+        Builder.createBranch(SEI->getLoc(), ThreadedSuccessorBlock,
+                             SEI->getOperand());
       }
     } else {
       Builder.createBranch(SEI->getLoc(), ThreadedSuccessorBlock,
@@ -3688,6 +3689,19 @@ bool SimplifyCFG::simplifyArgument(SILBasicBlock *BB, unsigned i) {
   auto *Use = *A->use_begin();
   auto *User = Use->getUser();
 
+  auto disableInOSSA = [](SingleValueInstruction *inst) {
+    assert(isa<StructInst>(inst) || isa<TupleInst>(inst) ||
+           isa<EnumInst>(inst));
+    if (!inst->getFunction()->hasOwnership()) {
+      return false;
+    }
+    if (inst->getOwnershipKind() == OwnershipKind::Owned)
+      return !inst->getSingleUse();
+    if (BorrowedValue borrow = BorrowedValue(inst->getOperand(0)))
+      return borrow.isLocalScope();
+    return false;
+  };
+
   // Handle projections.
   if (!isa<StructExtractInst>(User) &&
       !isa<TupleExtractInst>(User) &&
@@ -3702,15 +3716,17 @@ bool SimplifyCFG::simplifyArgument(SILBasicBlock *BB, unsigned i) {
       return false;
     auto *Branch = cast<BranchInst>(Pred->getTerminator());
     SILValue BranchArg = Branch->getArg(i);
-    if (isa<StructInst>(BranchArg))
-      continue;
-    if (isa<TupleInst>(BranchArg))
-      continue;
+    if (!isa<StructInst>(BranchArg) && !isa<TupleInst>(BranchArg) &&
+        !isa<EnumInst>(BranchArg)) {
+      return false;
+    }
     if (auto *EI = dyn_cast<EnumInst>(BranchArg)) {
-      if (EI->getElement() == cast<UncheckedEnumDataInst>(proj)->getElement())
-        continue;
+      if (EI->getElement() != cast<UncheckedEnumDataInst>(proj)->getElement())
+        return false;
+    }
+    if (disableInOSSA(cast<SingleValueInstruction>(BranchArg))) {
+      return false;
     }
-    return false;
   }
 
   // Okay, we'll replace the BB arg with one with the right type, replace

test/SILOptimizer/simplify_cfg_ossa_bbargs.sil

@@ -11,6 +11,10 @@ enum FakeOptional<T> {
   case none
 }
 
+struct NonTrivialStruct {
+  var val: Klass
+}
+
 sil @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
 
 // CHECK-LABEL: sil [ossa] @test_simplify_args1 :
@@ -172,7 +176,7 @@ bb4:
 }
 
 // CHECK-LABEL: sil [ossa] @test_simplify_args8 :
-// CHECK-NOT: unchecked_enum_date
+// CHECK-NOT: unchecked_enum_data
 // CHECK-LABEL: } // end sil function 'test_simplify_args8'
 sil [ossa] @test_simplify_args8 : $@convention(thin) (@guaranteed Klass) -> () {
 bb0(%0 : @guaranteed $Klass):
@@ -235,3 +239,77 @@ bb3(%4 : $Builtin.Int32):
   return %5 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_simplify_enum_arg_owned : 
+// CHECK: bb1([[ARG:%.*]] : @owned $FakeOptional<Klass>)
+// CHECK-LABEL: } // end sil function 'test_simplify_enum_arg_owned'
+sil [ossa] @test_simplify_enum_arg_owned : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  test_specification "simplify-cfg-simplify-argument @block[1] 0"
+  %1 = enum $FakeOptional<Klass>, #FakeOptional.some!enumelt, %0 : $Klass
+  %c = copy_value %1 : $FakeOptional<Klass>
+  br bb1(%1 : $FakeOptional<Klass>)
+
+bb1(%3 : @owned $FakeOptional<Klass>):
+  %4 = unchecked_enum_data %3 : $FakeOptional<Klass>, #FakeOptional.some!enumelt
+  %f = function_ref @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %f(%4) : $@convention(thin) (@guaranteed Klass) -> ()
+  br bb2
+
+bb2:
+  destroy_value %4 : $Klass
+  br bb3
+
+bb3:
+  destroy_value %c : $FakeOptional<Klass>
+  %t = tuple ()
+  return %t : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_simplify_enum_arg_guaranteed1 : 
+// CHECK: bb1([[ARG:%.*]] : @guaranteed $FakeOptional<Klass>)
+// CHECK-LABEL: } // end sil function 'test_simplify_enum_arg_guaranteed1'
+sil [ossa] @test_simplify_enum_arg_guaranteed1 : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  test_specification "simplify-cfg-simplify-argument @block[1] 0"
+  %b = begin_borrow %0 : $Klass
+  %1 = enum $FakeOptional<Klass>, #FakeOptional.some!enumelt, %b : $Klass
+  %c = copy_value %1 : $FakeOptional<Klass>
+  br bb1(%1 : $FakeOptional<Klass>)
+
+bb1(%3 : @guaranteed $FakeOptional<Klass>):
+  %4 = unchecked_enum_data %3 : $FakeOptional<Klass>, #FakeOptional.some!enumelt
+  %f = function_ref @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %f(%4) : $@convention(thin) (@guaranteed Klass) -> ()
+  br bb2
+
+bb2:
+  destroy_value %c : $FakeOptional<Klass>
+  end_borrow %b : $Klass
+  destroy_value %0 : $Klass
+  %t = tuple ()
+  return %t : $()
+}
+
+// CHECK: bb1([[ARG:%.*]] : @guaranteed $NonTrivialStruct)
+sil [ossa] @test_simplify_enum_arg_guaranteed2 : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  test_specification "simplify-cfg-simplify-argument @block[1] 0"
+  %b = begin_borrow %0 : $Klass
+  %1 = struct $NonTrivialStruct(%b : $Klass) 
+  %c = copy_value %1 : $NonTrivialStruct
+  br bb1(%1 : $NonTrivialStruct)
+
+bb1(%3 : @guaranteed $NonTrivialStruct):
+  %4 = struct_extract %3 : $NonTrivialStruct, #NonTrivialStruct.val
+  %f = function_ref @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %f(%4) : $@convention(thin) (@guaranteed Klass) -> ()
+  br bb2
+
+bb2:
+  destroy_value %c : $NonTrivialStruct
+  end_borrow %b : $Klass
+  destroy_value %0 : $Klass
+  %t = tuple ()
+  return %t : $()
+}
+