SimplifyBeginCOWMutation: don't eliminate end_cow_mutation instructions with the [keep_unique] flag

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBeginCOWMutation.swift

@@ -74,7 +74,13 @@ private extension BeginCOWMutationInst {
       return false
     }
     let buffer = instanceResult
-    if buffer.uses.ignoreDebugUses.contains(where: { !($0.instruction is EndCOWMutationInst) }) {
+    guard buffer.uses.ignoreDebugUses.allSatisfy({
+        if let endCOW = $0.instruction as? EndCOWMutationInst {
+          return !endCOW.doKeepUnique
+        }
+        return false
+      }) else
+    {
       return false
     }
 
@@ -91,7 +97,8 @@ private extension BeginCOWMutationInst {
     if !uniquenessResult.uses.ignoreDebugUses.isEmpty {
       return false
     }
-    guard let endCOW = instance as? EndCOWMutationInst else {
+    guard let endCOW = instance as? EndCOWMutationInst,
+          !endCOW.doKeepUnique else {
       return false
     }
     if endCOW.uses.ignoreDebugUses.contains(where: { $0.instruction != self }) {

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -592,7 +592,8 @@ SILInstruction *SILCombiner::visitEndCOWMutationInst(EndCOWMutationInst *ECM) {
 
   SingleValueInstruction *refCast = cast<SingleValueInstruction>(op);
   auto *newECM = Builder.createEndCOWMutation(ECM->getLoc(),
-                                              refCast->getOperand(0));
+                                              refCast->getOperand(0),
+                                              ECM->doKeepUnique());
   ECM->replaceAllUsesWith(refCast);
   refCast->setOperand(0, newECM);
   refCast->moveAfter(newECM);

test/SILOptimizer/reverse-array.swift

@@ -36,7 +36,7 @@ public func reverseArray(_ a: [Int]) -> [Int] {
 // CHECK:     begin_cow_mutation
 // CHECK-NOT: {{.*(_cow_mutation|cond_fail)}}
 // CHECK:     end_cow_mutation
-// CHECK-NOT: {{.*(_cow_mutation|cond_fail)}}
+// CHECK:     end_cow_mutation
 
 // In SIL we fail to eliminate the bounds check of the input array.
 // But that's okay, because LLVM can do that.

test/SILOptimizer/sil_combine_inst_passes.sil

@@ -48,6 +48,17 @@ bb0(%0 : @owned $Buffer):
   return %t : $(Builtin.Int1, Buffer)
 }
 
+// CHECK-LABEL: sil [ossa] @dont_remove_end_begin_cow_pair3
+// CHECK:   end_cow_mutation
+// CHECK:   begin_cow_mutation
+// CHECK: } // end sil function 'dont_remove_end_begin_cow_pair3'
+sil [ossa] @dont_remove_end_begin_cow_pair3 : $@convention(thin) (@owned Buffer) -> @owned Buffer {
+bb0(%0 : @owned $Buffer):
+  %e = end_cow_mutation [keep_unique] %0 : $Buffer
+  (%u, %b) = begin_cow_mutation %e : $Buffer
+  return %b : $Buffer
+}
+
 // CHECK-LABEL: sil [ossa] @remove_begin_end_cow_pair
 // CHECK-NOT: end_cow_mutation
 // CHECK-NOT: begin_cow_mutation
@@ -87,6 +98,16 @@ bb0(%0 : @owned $Buffer):
   return %t : $(Builtin.Int1, Buffer)
 }
 
+// CHECK-LABEL: sil [ossa] @dont_remove_begin_end_cow_pair3
+// CHECK:   begin_cow_mutation
+// CHECK:   end_cow_mutation
+// CHECK: } // end sil function 'dont_remove_begin_end_cow_pair3'
+sil [ossa] @dont_remove_begin_end_cow_pair3 : $@convention(thin) (@owned Buffer) -> @owned Buffer {
+bb0(%0 : @owned $Buffer):
+  (%u, %b) = begin_cow_mutation %0 : $Buffer
+  %e = end_cow_mutation [keep_unique] %b : $Buffer
+  return %e : $Buffer
+}
 // CHECK-LABEL: sil @optimize_empty_cow_singleton
 // CHECK:   [[I:%[0-9]+]] = integer_literal $Builtin.Int1, 0
 // CHECK:   begin_cow_mutation

test/SILOptimizer/sil_combine_ossa.sil

@@ -1067,6 +1067,30 @@ bb0(%0 : @owned $C1):
   return %9999 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @end_cow_mutation_of_cast :
+// CHECK:         %1 = end_cow_mutation %0
+// CHECK-NEXT:    %2 = upcast %1
+// CHECK-NEXT:    return %2
+// CHECK: } // end sil function 'end_cow_mutation_of_cast'
+sil [ossa] @end_cow_mutation_of_cast : $@convention(thin) (@owned C2) -> @owned C1 {
+bb0(%0 : @owned $C2):
+  %1 = upcast %0 : $C2 to $C1
+  %2 = end_cow_mutation %1 : $C1
+  return %2 : $C1
+}
+
+// CHECK-LABEL: sil [ossa] @end_cow_mutation_of_cast2 :
+// CHECK:         %1 = end_cow_mutation [keep_unique] %0
+// CHECK-NEXT:    %2 = upcast %1
+// CHECK-NEXT:    return %2
+// CHECK: } // end sil function 'end_cow_mutation_of_cast2'
+sil [ossa] @end_cow_mutation_of_cast2 : $@convention(thin) (@owned C2) -> @owned C1 {
+bb0(%0 : @owned $C2):
+  %1 = upcast %0 : $C2 to $C1
+  %2 = end_cow_mutation [keep_unique] %1 : $C1
+  return %2 : $C1
+}
+
 struct XS {
   var m: Int
   var k: Float