Fixes to utilities used by OSSALifetimeCompletion

include/swift/SIL/OwnershipUseVisitor.h

@@ -283,11 +283,18 @@ bool OwnershipUseVisitor<Impl>::visitInnerBorrowScopeEnd(Operand *borrowEnd) {
   case OperandOwnership::EndBorrow:
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
 
-  case OperandOwnership::Reborrow:
+  case OperandOwnership::Reborrow: {
     if (!asImpl().handleInnerReborrow(borrowEnd))
       return false;
 
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
+  }
+  case OperandOwnership::DestroyingConsume: {
+    // 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());
+    return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
+  }
 
   default:
     llvm_unreachable("expected borrow scope end");

include/swift/SIL/OwnershipUtils.h

@@ -475,6 +475,9 @@ struct BorrowingOperand {
   /// valid BorrowedValue instance.
   BorrowedValue getBorrowIntroducingUserResult();
 
+  /// Return the borrowing operand's value.
+  SILValue getScopeIntroducingUserResult();
+
   /// Compute the implicit uses that this borrowing operand "injects" into the
   /// set of its operands uses.
   ///

lib/SIL/Utils/OwnershipLiveness.cpp

@@ -140,8 +140,10 @@ struct InteriorLivenessVisitor :
   /// Handles begin_borrow, load_borrow, store_borrow, begin_apply.
   bool handleInnerBorrow(BorrowingOperand borrowingOperand) {
     if (handleInnerScopeCallback) {
-      handleInnerScopeCallback(
-        borrowingOperand.getBorrowIntroducingUserResult().value);
+      auto value = borrowingOperand.getScopeIntroducingUserResult();
+      if (value) {
+        handleInnerScopeCallback(value);
+      }
     }
     return true;
   }

lib/SIL/Utils/OwnershipUtils.cpp

@@ -751,6 +751,30 @@ BorrowedValue BorrowingOperand::getBorrowIntroducingUserResult() {
   llvm_unreachable("covered switch");
 }
 
+SILValue BorrowingOperand::getScopeIntroducingUserResult() {
+  switch (kind) {
+  case BorrowingOperandKind::Invalid:
+  case BorrowingOperandKind::Yield:
+  case BorrowingOperandKind::Apply:
+  case BorrowingOperandKind::TryApply:
+    return SILValue();
+
+  case BorrowingOperandKind::BeginAsyncLet:
+  case BorrowingOperandKind::PartialApplyStack:
+  case BorrowingOperandKind::BeginBorrow:
+    return cast<SingleValueInstruction>(op->getUser());
+
+  case BorrowingOperandKind::BeginApply:
+    return cast<BeginApplyInst>(op->getUser())->getTokenResult();
+
+  case BorrowingOperandKind::Branch: {
+    PhiOperand phiOp(op);
+    return phiOp.getValue();
+  }
+  }
+  llvm_unreachable("covered switch");
+}
+
 void BorrowingOperand::getImplicitUses(
     SmallVectorImpl<Operand *> &foundUses) const {
   // FIXME: this visitScopeEndingUses should never return false once dead

test/SILOptimizer/ossa_lifetime_completion.sil

@@ -102,6 +102,31 @@ bb3:
   return %r : $()
 }
 
+sil @use_guaranteed : $@convention(thin) (@guaranteed C) -> ()
+
+sil [ossa] @argTest : $@convention(method) (@owned C) -> () {
+bb0(%0 : @owned $C):
+  test_specification "ossa-lifetime-completion @argument"
+  debug_value %0 : $C
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb4
+
+bb2:
+  br bb3
+
+bb3:
+  %3 = function_ref @use_guaranteed : $@convention(thin) (@guaranteed C) -> ()
+  %4 = apply %3(%0) : $@convention(thin) (@guaranteed C) -> ()
+  destroy_value %0 : $C
+  %r = tuple ()
+  return %r : $()
+
+bb4:
+  unreachable
+}
+
 // Ensure no assert fires while inserting dead end blocks to the worklist
 sil [ossa] @testLexicalLifetimeCompletion : $@convention(thin) (@owned C) -> () {
 bb0(%0 : @owned $C):
@@ -143,3 +168,29 @@ bb10:
   br bb8
 }
 
+sil @foo : $@convention(thin) (@guaranteed C) -> ()
+
+// Ensure no assert fires while handling lifetime end of partial_apply
+sil [ossa] @testPartialApplyStack : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  test_specification "ossa-lifetime-completion @instruction[0]"
+  %8 = copy_value %0 : $C
+  %9 = begin_borrow %8 : $C
+  %80 = function_ref @foo : $@convention(thin) (@guaranteed C) -> ()
+  %81 = partial_apply [callee_guaranteed] [on_stack] %80(%9) : $@convention(thin) (@guaranteed C) -> ()
+  cond_br undef, bb1, bb2
+
+bb1:
+  destroy_value %81 : $@noescape @callee_guaranteed () -> ()
+  br bb3
+
+bb2:
+  unreachable
+
+bb3:
+  end_borrow %9 : $C
+  destroy_value %8 : $C
+  %180 = tuple ()
+  return %180 : $()
+}
+