[LifetimeCompletion] Complete coro token lifetimes with end_borrow.

include/swift/SIL/OwnershipUtils.h

@@ -490,6 +490,7 @@ class BorrowedValueKind {
     BeginBorrow,
     SILFunctionArgument,
     Phi,
+    BeginApplyToken,
   };
 
 private:
@@ -520,6 +521,13 @@ class BorrowedValueKind {
       }
       return Kind::Phi;
     }
+    case ValueKind::MultipleValueInstructionResult:
+      if (!isaResultOf<BeginApplyInst>(value))
+        return Kind::Invalid;
+      if (value->isBeginApplyToken()) {
+        return Kind::BeginApplyToken;
+      }
+      return Kind::Invalid;
     }
   }
 
@@ -540,6 +548,7 @@ class BorrowedValueKind {
     case BorrowedValueKind::BeginBorrow:
     case BorrowedValueKind::LoadBorrow:
     case BorrowedValueKind::Phi:
+    case BorrowedValueKind::BeginApplyToken:
       return true;
     case BorrowedValueKind::SILFunctionArgument:
       return false;

include/swift/SIL/SILInstruction.h

@@ -3163,6 +3163,7 @@ class PartialApplyInst final
 
 class EndApplyInst;
 class AbortApplyInst;
+class EndBorrowInst;
 
 /// BeginApplyInst - Represents the beginning of the full application of
 /// a yield_once coroutine (up until the coroutine yields a value back).
@@ -3216,10 +3217,13 @@ class BeginApplyInst final
 
   void getCoroutineEndPoints(
       SmallVectorImpl<EndApplyInst *> &endApplyInsts,
-      SmallVectorImpl<AbortApplyInst *> &abortApplyInsts) const;
+      SmallVectorImpl<AbortApplyInst *> &abortApplyInsts,
+      SmallVectorImpl<EndBorrowInst *> *endBorrowInsts = nullptr) const;
 
-  void getCoroutineEndPoints(SmallVectorImpl<Operand *> &endApplyInsts,
-                             SmallVectorImpl<Operand *> &abortApplyInsts) const;
+  void getCoroutineEndPoints(
+      SmallVectorImpl<Operand *> &endApplyInsts,
+      SmallVectorImpl<Operand *> &abortApplyInsts,
+      SmallVectorImpl<Operand *> *endBorrowInsts = nullptr) const;
 };
 
 /// AbortApplyInst - Unwind the full application of a yield_once coroutine.
@@ -4539,8 +4543,6 @@ class StoreInst
   }
 };
 
-class EndBorrowInst;
-
 /// Represents a load of a borrowed value. Must be paired with an end_borrow
 /// instruction in its use-def list.
 class LoadBorrowInst :

include/swift/SIL/SILValue.h

@@ -590,6 +590,8 @@ class ValueBase : public SILNode, public SILAllocated<ValueBase> {
 
   bool isGuaranteedForwarding() const;
 
+  bool isBeginApplyToken() const;
+
   /// Unsafely eliminate moveonly from this value's type. Returns true if the
   /// value's underlying type was move only and thus was changed. Returns false
   /// otherwise.

lib/SIL/IR/SILInstructions.cpp

@@ -693,29 +693,41 @@ BeginApplyInst *BeginApplyInst::create(
 
 void BeginApplyInst::getCoroutineEndPoints(
     SmallVectorImpl<EndApplyInst *> &endApplyInsts,
-    SmallVectorImpl<AbortApplyInst *> &abortApplyInsts) const {
+    SmallVectorImpl<AbortApplyInst *> &abortApplyInsts,
+    SmallVectorImpl<EndBorrowInst *> *endBorrowInsts) const {
   for (auto *tokenUse : getTokenResult()->getUses()) {
     auto *user = tokenUse->getUser();
     if (auto *end = dyn_cast<EndApplyInst>(user)) {
       endApplyInsts.push_back(end);
       continue;
     }
-
-    abortApplyInsts.push_back(cast<AbortApplyInst>(user));
+    if (auto *abort = dyn_cast<AbortApplyInst>(user)) {
+      abortApplyInsts.push_back(abort);
+      continue;
+    }
+    auto *end = cast<EndBorrowInst>(user);
+    if (endBorrowInsts) {
+      endBorrowInsts->push_back(end);
+    }
   }
 }
 
 void BeginApplyInst::getCoroutineEndPoints(
     SmallVectorImpl<Operand *> &endApplyInsts,
-    SmallVectorImpl<Operand *> &abortApplyInsts) const {
+    SmallVectorImpl<Operand *> &abortApplyInsts,
+    SmallVectorImpl<Operand *> *endBorrowInsts) const {
   for (auto *tokenUse : getTokenResult()->getUses()) {
     auto *user = tokenUse->getUser();
     if (isa<EndApplyInst>(user)) {
       endApplyInsts.push_back(tokenUse);
       continue;
     }
+    if (isa<AbortApplyInst>(user)) {
+      abortApplyInsts.push_back(tokenUse);
+      continue;
+    }
 
-    assert(isa<AbortApplyInst>(user));
+    assert(isa<EndBorrowInst>(user));
     abortApplyInsts.push_back(tokenUse);
   }
 }

lib/SIL/IR/SILValue.cpp

@@ -191,6 +191,13 @@ bool ValueBase::isGuaranteedForwarding() const {
   return phi->isGuaranteedForwarding();
 }
 
+bool ValueBase::isBeginApplyToken() const {
+  auto *result = isaResultOf<BeginApplyInst>(this);
+  if (!result)
+    return false;
+  return result->isBeginApplyToken();
+}
+
 bool ValueBase::hasDebugTrace() const {
   for (auto *op : getUses()) {
     if (auto *debugValue = dyn_cast<DebugValueInst>(op->getUser())) {

lib/SIL/Utils/OwnershipUtils.cpp

@@ -913,6 +913,9 @@ void BorrowedValueKind::print(llvm::raw_ostream &os) const {
   case BorrowedValueKind::Phi:
     os << "Phi";
     return;
+  case BorrowedValueKind::BeginApplyToken:
+    os << "BeginApply";
+    return;
   }
   llvm_unreachable("Covered switch isn't covered?!");
 }
@@ -945,6 +948,7 @@ bool BorrowedValue::visitLocalScopeEndingUses(
   case BorrowedValueKind::LoadBorrow:
   case BorrowedValueKind::BeginBorrow:
   case BorrowedValueKind::Phi:
+  case BorrowedValueKind::BeginApplyToken:
     for (auto *use : lookThroughBorrowedFromUser(value)->getUses()) {
       if (use->isLifetimeEnding()) {
         if (!visitor(use))
@@ -1819,6 +1823,7 @@ class FindEnclosingDefs {
 
       case BorrowedValueKind::LoadBorrow:
       case BorrowedValueKind::SILFunctionArgument:
+      case BorrowedValueKind::BeginApplyToken:
         // There is no enclosing def on this path.
         return true;
       }
@@ -2037,6 +2042,7 @@ class FindEnclosingDefs {
     }
     case BorrowedValueKind::LoadBorrow:
     case BorrowedValueKind::SILFunctionArgument:
+    case BorrowedValueKind::BeginApplyToken:
       // There is no enclosing def on this path.
       break;
     }

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -447,7 +447,8 @@ static bool visitScopeEndsRequiringInit(
   if (auto bai =
           dyn_cast_or_null<BeginApplyInst>(operand->getDefiningInstruction())) {
     for (auto *inst : bai->getTokenResult()->getUsers()) {
-      assert(isa<EndApplyInst>(inst) || isa<AbortApplyInst>(inst));
+      assert(isa<EndApplyInst>(inst) || isa<AbortApplyInst>(inst) ||
+             isa<EndBorrowInst>(inst));
       visit(inst, ScopeRequiringFinalInit::Coroutine);
     }
     return true;

lib/SILOptimizer/Utils/CanonicalizeBorrowScope.cpp

@@ -140,6 +140,7 @@ SILValue CanonicalizeBorrowScope::getCanonicalBorrowedDef(SILValue def) {
 
       case BorrowedValueKind::LoadBorrow:
       case BorrowedValueKind::Phi:
+      case BorrowedValueKind::BeginApplyToken:
         break;
       }
     }
@@ -170,6 +171,7 @@ bool CanonicalizeBorrowScope::computeBorrowLiveness() {
     // can handle persistentCopies.
     return false;
   case BorrowedValueKind::BeginBorrow:
+  case BorrowedValueKind::BeginApplyToken:
     break;
   }
   // Note that there is no need to look through any reborrows. The reborrowed

lib/SILOptimizer/Utils/LoopUtils.cpp

@@ -307,7 +307,8 @@ bool swift::canDuplicateLoopInstruction(SILLoop *L, SILInstruction *I) {
   if (auto BAI = dyn_cast<BeginApplyInst>(I)) {
     for (auto UI : BAI->getTokenResult()->getUses()) {
       auto User = UI->getUser();
-      assert(isa<EndApplyInst>(User) || isa<AbortApplyInst>(User));
+      assert(isa<EndApplyInst>(User) || isa<AbortApplyInst>(User) ||
+             isa<EndBorrowInst>(User));
       if (!L->contains(User))
         return false;
     }

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -41,10 +41,11 @@ static bool canInlineBeginApply(BeginApplyInst *BA) {
     if (isa<EndApplyInst>(user)) {
       if (hasEndApply) return false;
       hasEndApply = true;
-    } else {
-      assert(isa<AbortApplyInst>(user));
+    } else if (isa<AbortApplyInst>(user)) {
       if (hasAbortApply) return false;
       hasAbortApply = true;
+    } else {
+      assert(isa<EndBorrowInst>(user));
     }
   }
 
@@ -95,6 +96,8 @@ class BeginApplySite {
   SILBasicBlock *AbortApplyBB = nullptr;
   SILBasicBlock *AbortApplyReturnBB = nullptr;
 
+  SmallVector<EndBorrowInst *, 2> EndBorrows;
+
 public:
   BeginApplySite(BeginApplyInst *BeginApply, SILLocation Loc,
                  SILBuilder *Builder)
@@ -112,7 +115,8 @@ class BeginApplySite {
                   SmallVectorImpl<SILInstruction *> &endBorrowInsertPts) {
     SmallVector<EndApplyInst *, 1> endApplyInsts;
     SmallVector<AbortApplyInst *, 1> abortApplyInsts;
-    BeginApply->getCoroutineEndPoints(endApplyInsts, abortApplyInsts);
+    BeginApply->getCoroutineEndPoints(endApplyInsts, abortApplyInsts,
+                                      &EndBorrows);
     while (!endApplyInsts.empty()) {
       auto *endApply = endApplyInsts.pop_back_val();
       collectEndApply(endApply);
@@ -257,6 +261,8 @@ class BeginApplySite {
       EndApply->eraseFromParent();
     if (AbortApply)
       AbortApply->eraseFromParent();
+    for (auto *EndBorrow : EndBorrows)
+      EndBorrow->eraseFromParent();
 
     assert(!BeginApply->hasUsesOfAnyResult());
   }

test/SILOptimizer/inline_begin_apply.sil

@@ -690,3 +690,118 @@ bb_unwind:
   unwind
 
 }
+
+sil [ossa] [transparent] @simplest_in_guaranteed : $@yield_once @convention(thin) () -> @yields @in_guaranteed () {
+entry:
+  %tuple = tuple ()
+  %addr = alloc_stack $()
+  store %tuple to [trivial] %addr : $*()
+  yield %addr : $*(), resume resume_block, unwind unwind_block
+
+resume_block:
+  dealloc_stack %addr : $*()
+  %retval = tuple ()
+  return %retval : $()
+
+unwind_block:
+  dealloc_stack %addr : $*()
+  unwind
+}
+
+// CHECK-LABEL: sil [ossa] @ends_1 : {{.*}} {
+// CHECK:         [[VOID:%[^,]+]] = tuple ()
+// CHECK:         [[ADDR:%[^,]+]] = alloc_stack $()
+// CHECK:         store [[VOID]] to [trivial] [[ADDR]]
+// CHECK:         cond_br undef, [[GOOD:bb[0-9]+]], [[BAD:bb[0-9]+]]
+// CHECK:       [[GOOD]]:
+// CHECK:         dealloc_stack [[ADDR]]
+// CHECK:         br [[EXIT:bb[0-9]+]]
+// CHECK:       [[BAD]]:
+// CHECK:         cond_br undef, [[NORMAL_BAD:bb[0-9]+]], [[REALLY_BAD:bb[0-9]+]]
+// CHECK:       [[NORMAL_BAD]]:
+// CHECK:         dealloc_stack [[ADDR]]
+// CHECK:         br [[EXIT]]
+// CHECK:       [[REALLY_BAD]]:
+// CHECK:         unreachable
+// CHECK:       [[EXIT]]:
+// CHECK-LABEL: } // end sil function 'ends_1'
+sil [ossa] @ends_1 : $@convention(thin) () -> () {
+entry:
+  %simplest = function_ref @simplest_in_guaranteed : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
+  (%_, %token) = begin_apply %simplest() : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
+  cond_br undef, good, bad
+
+good:
+  end_apply %token as $()
+  br exit
+
+bad:
+  cond_br undef, normal_bad, really_bad
+
+normal_bad:
+  abort_apply %token
+  br exit
+
+really_bad:
+  end_borrow %token : $Builtin.SILToken
+  unreachable
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// Test handling for multiple end_borrows of a token.
+// CHECK-LABEL: sil [ossa] @ends_2 : {{.*}} {
+// CHECK:         [[VOID:%[^,]+]] = tuple ()
+// CHECK:         [[ADDR:%[^,]+]] = alloc_stack $()
+// CHECK:         store [[VOID]] to [trivial] [[ADDR]] : $*()
+// CHECK:         cond_br undef, [[GOOD:bb[0-9]+]], [[BAD:bb[0-9]+]]
+// CHECK:       [[GOOD]]:
+// CHECK:         dealloc_stack [[ADDR]]
+// CHECK:         br [[EXIT:bb[0-9]+]]
+// CHECK:       [[BAD]]:
+// CHECK:         cond_br undef, [[NORMAL_BAD:bb[0-9]+]], [[REALLY_BAD:bb[0-9]+]]
+// CHECK:       [[NORMAL_BAD]]:
+// CHECK:         dealloc_stack [[ADDR]]
+// CHECK:         br [[EXIT]]
+// CHECK:       [[REALLY_BAD]]:
+// CHECK:         cond_br undef, [[BAD_NEWS_BUNNIES:bb[0-9]+]], [[BAD_NEWS_BEARS:bb[0-9]+]]
+// CHECK:       [[BAD_NEWS_BUNNIES]]:
+// CHECK:         unreachable
+// CHECK:       [[BAD_NEWS_BEARS]]:
+// CHECK:         unreachable
+// CHECK:       [[EXIT]]:
+// CHECK-LABEL: } // end sil function 'ends_2'
+sil [ossa] @ends_2 : $@convention(thin) () -> () {
+entry:
+  %simplest = function_ref @simplest_in_guaranteed : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
+  (%_, %token) = begin_apply %simplest() : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
+  cond_br undef, good, bad
+
+good:
+  end_apply %token as $()
+  br exit
+
+bad:
+  cond_br undef, normal_bad, really_bad
+
+normal_bad:
+  abort_apply %token
+  br exit
+
+really_bad:
+  cond_br undef, bad_news_bunnies, bad_news_bears
+
+bad_news_bunnies:
+  end_borrow %token : $Builtin.SILToken
+  unreachable
+
+bad_news_bears:
+  end_borrow %token : $Builtin.SILToken
+  unreachable
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}

test/SILOptimizer/ossa_lifetime_completion.sil

@@ -513,3 +513,27 @@ entry(%instance : @owned $C):
   store %instance to [init] %addr : $*C
   unreachable
 }
+
+// CHECK-LABEL: sil [ossa] @begin_apply : {{.*}} {
+// CHECK:         ({{%[^,]+}}, [[TOKEN:%[^,]+]]) = begin_apply undef()
+// CHECK:         cond_br undef, [[LEFT:bb[0-9]+]], [[RIGHT:bb[0-9]+]]
+// CHECK:       [[LEFT]]:
+// CHECK:         end_apply [[TOKEN]]
+// CHECK:       [[RIGHT]]:
+// CHECK:         end_borrow [[TOKEN]]
+// CHECK:         unreachable
+// CHECK-LABEL: } // end sil function 'begin_apply'
+sil [ossa] @begin_apply : $@convention(thin) () -> () {
+entry:
+  (%_, %token) = begin_apply undef() : $@yield_once @convention(thin) () -> (@yields @in_guaranteed ())
+  specify_test "ossa-lifetime-completion %token"
+  cond_br undef, left, right
+
+left:
+  end_apply %token as $()
+  %retval = tuple ()
+  return %retval : $()
+
+right:
+  unreachable
+}