[CoroutineAccessors] Fix begin_apply for protos.

When the CoroutineAccessors feature is enabled, `begin_apply`
instructions produce an additional result representing the allocation
done by the callee.  Fix a couple of cases where this additional result
was not being handled.

Author: nate-chandler

lib/SILGen/SILGenApply.cpp

@@ -6000,7 +6000,8 @@ SILValue SILGenFunction::emitApplyWithRethrow(SILLocation loc, SILValue fn,
   return normalBB->createPhiArgument(resultType, OwnershipKind::Owned);
 }
 
-std::pair<MultipleValueInstructionResult *, CleanupHandle>
+std::tuple<MultipleValueInstructionResult *, CleanupHandle, SILValue,
+           CleanupHandle>
 SILGenFunction::emitBeginApplyWithRethrow(SILLocation loc, SILValue fn,
                                           SILType substFnType,
                                           SubstitutionMap subs,
@@ -6021,19 +6022,29 @@ SILGenFunction::emitBeginApplyWithRethrow(SILLocation loc, SILValue fn,
 
   auto *token = beginApply->getTokenResult();
 
+  CleanupHandle deallocCleanup;
+  auto *allocation = beginApply->getCalleeAllocationResult();
+  if (allocation) {
+    deallocCleanup = enterDeallocStackCleanup(allocation);
+  }
+
   Cleanups.pushCleanup<EndCoroutineApply>(token);
   auto abortCleanup = Cleanups.getTopCleanup();
 
-  return { token, abortCleanup };
+  return {token, abortCleanup, allocation, deallocCleanup};
 }
 
-void SILGenFunction::emitEndApplyWithRethrow(SILLocation loc,
-                                        MultipleValueInstructionResult *token) {
+void SILGenFunction::emitEndApplyWithRethrow(
+    SILLocation loc, MultipleValueInstructionResult *token,
+    SILValue allocation) {
   // TODO: adjust this to handle results of TryBeginApplyInst.
   assert(token->isBeginApplyToken());
 
   B.createEndApply(loc, token,
                    SILType::getEmptyTupleType(getASTContext()));
+  if (allocation) {
+    B.createDeallocStack(loc, allocation);
+  }
 }
 
 void SILGenFunction::emitYield(SILLocation loc,

lib/SILGen/SILGenFunction.h

@@ -2185,12 +2185,14 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                                 SubstitutionMap subs,
                                 ArrayRef<SILValue> args);
 
-  std::pair<MultipleValueInstructionResult *, CleanupHandle>
+  std::tuple<MultipleValueInstructionResult *, CleanupHandle, SILValue,
+             CleanupHandle>
   emitBeginApplyWithRethrow(SILLocation loc, SILValue fn, SILType substFnType,
                             SubstitutionMap subs, ArrayRef<SILValue> args,
                             SmallVectorImpl<SILValue> &yields);
   void emitEndApplyWithRethrow(SILLocation loc,
-                               MultipleValueInstructionResult *token);
+                               MultipleValueInstructionResult *token,
+                               SILValue allocation);
 
   ManagedValue emitExtractFunctionIsolation(SILLocation loc,
                                         ArgumentSource &&fnValue);

lib/SILGen/SILGenPoly.cpp

@@ -6824,10 +6824,13 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
   case SILCoroutineKind::YieldOnce:
   case SILCoroutineKind::YieldOnce2: {
     SmallVector<SILValue, 4> derivedYields;
-    auto tokenAndCleanup =
-        emitBeginApplyWithRethrow(loc, derivedRef,
-                                  SILType::getPrimitiveObjectType(derivedFTy),
-                                  subs, args, derivedYields);
+    auto tokenAndCleanups = emitBeginApplyWithRethrow(
+        loc, derivedRef, SILType::getPrimitiveObjectType(derivedFTy), subs,
+        args, derivedYields);
+    auto token = std::get<0>(tokenAndCleanups);
+    auto abortCleanup = std::get<1>(tokenAndCleanups);
+    auto allocation = std::get<2>(tokenAndCleanups);
+    auto deallocCleanup = std::get<3>(tokenAndCleanups);
     auto overrideSubs = SubstitutionMap::getOverrideSubstitutions(
         base.getDecl(), derived.getDecl()).subst(subs);
 
@@ -6837,10 +6840,13 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
     translateYields(*this, loc, derivedYields, derivedYieldInfo, baseYieldInfo);
 
     // Kill the abort cleanup without emitting it.
-    Cleanups.setCleanupState(tokenAndCleanup.second, CleanupState::Dead);
+    Cleanups.setCleanupState(abortCleanup, CleanupState::Dead);
+    if (allocation) {
+      Cleanups.setCleanupState(deallocCleanup, CleanupState::Dead);
+    }
 
     // End the inner coroutine normally.
-    emitEndApplyWithRethrow(loc, tokenAndCleanup.first);
+    emitEndApplyWithRethrow(loc, token, allocation);
 
     result = B.createTuple(loc, {});
     break;
@@ -7213,9 +7219,12 @@ void SILGenFunction::emitProtocolWitness(
   case SILCoroutineKind::YieldOnce:
   case SILCoroutineKind::YieldOnce2: {
     SmallVector<SILValue, 4> witnessYields;
-    auto tokenAndCleanup =
-      emitBeginApplyWithRethrow(loc, witnessFnRef, witnessSILTy, witnessSubs,
-                                args, witnessYields);
+    auto tokenAndCleanups = emitBeginApplyWithRethrow(
+        loc, witnessFnRef, witnessSILTy, witnessSubs, args, witnessYields);
+    auto token = std::get<0>(tokenAndCleanups);
+    auto abortCleanup = std::get<1>(tokenAndCleanups);
+    auto allocation = std::get<2>(tokenAndCleanups);
+    auto deallocCleanup = std::get<3>(tokenAndCleanups);
 
     YieldInfo witnessYieldInfo(SGM, witness, witnessFTy, witnessSubs);
     YieldInfo reqtYieldInfo(SGM, requirement, thunkTy,
@@ -7224,10 +7233,13 @@ void SILGenFunction::emitProtocolWitness(
     translateYields(*this, loc, witnessYields, witnessYieldInfo, reqtYieldInfo);
 
     // Kill the abort cleanup without emitting it.
-    Cleanups.setCleanupState(tokenAndCleanup.second, CleanupState::Dead);
+    Cleanups.setCleanupState(abortCleanup, CleanupState::Dead);
+    if (allocation) {
+      Cleanups.setCleanupState(deallocCleanup, CleanupState::Dead);
+    }
 
     // End the inner coroutine normally.
-    emitEndApplyWithRethrow(loc, tokenAndCleanup.first);
+    emitEndApplyWithRethrow(loc, token, allocation);
 
     reqtResultValue = B.createTuple(loc, {});
     break;