Update assertion for partial_apply [on_stack] lifetime ends in lifetime completion

Since partial_apply [on_stack] is considered as a BorrowingOperand now,
and BorrowingOperand::visitScopeEndingUses can be called on the its forwarded values
lifetime ends, update the assertion to reflect this.

Author: meg-gupta

include/swift/SIL/OwnershipUseVisitor.h

@@ -290,9 +290,14 @@ bool OwnershipUseVisitor<Impl>::visitInnerBorrowScopeEnd(Operand *borrowEnd) {
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
   }
   case OperandOwnership::DestroyingConsume: {
-    // partial_apply [on_stack] can introduce borrowing operand and can have destroy_value consumes.
+    // partial_apply [on_stack] can introduce borrowing operand and can have
+    // destroy_value consumes.
     auto *pai = dyn_cast<PartialApplyInst>(borrowEnd->get());
-    assert(pai && pai->isOnStack());
+    // TODO: When we have ForwardingInstruction abstraction, walk the use-def
+    // chain to ensure we have a partial_apply [on_stack] def.
+    assert(pai && pai->isOnStack() ||
+           OwnershipForwardingMixin::get(
+               cast<SingleValueInstruction>(borrowEnd->get())));
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
   }
 

test/SILOptimizer/ossa_lifetime_completion.sil

@@ -171,7 +171,7 @@ bb10:
 sil @foo : $@convention(thin) (@guaranteed C) -> ()
 
 // Ensure no assert fires while handling lifetime end of partial_apply
-sil [ossa] @testPartialApplyStack : $@convention(thin) (@guaranteed C) -> () {
+sil [ossa] @testPartialApplyStack1 : $@convention(thin) (@guaranteed C) -> () {
 bb0(%0 : @guaranteed $C):
   test_specification "ossa-lifetime-completion @instruction[0]"
   %8 = copy_value %0 : $C
@@ -194,3 +194,21 @@ bb3:
   return %180 : $()
 }
 
+// Ensure no assert fires while handling lifetime end of partial_apply
+sil [ossa] @testPartialApplyStack2 : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  test_specification "ossa-lifetime-completion @instruction[1]"
+  %2 = alloc_stack $C
+  %3 = copy_value %0 : $C
+  %4 = begin_borrow %3 : $C
+  %5 = function_ref @foo : $@convention(thin) (@guaranteed C) -> ()
+  %6 = partial_apply [callee_guaranteed] [on_stack] %5(%4) : $@convention(thin) (@guaranteed C) -> ()
+  %7 = mark_dependence %6 : $@noescape @callee_guaranteed () -> () on %2 : $*C
+  destroy_value %7 : $@noescape @callee_guaranteed () -> ()
+  end_borrow %4 : $C
+  destroy_value %3 : $C
+  dealloc_stack %2 : $*C
+  %12 = tuple ()
+  return %12 : $()
+}
+