[CoroutineAccessors] SIL represents callee alloc.

When its operand has coroutine kind `yield_once_2`, a `begin_apply`
instruction produces an additional value representing the storage
allocated by the callee.  This storage must be deallocated by a
`dealloc_stack` on every path out of the function.  Like any other stack
allocation, it must obey stack discipline.

Author: nate-chandler

docs/SIL.rst

@@ -6132,23 +6132,36 @@ begin_apply
   // %float : $Float
   // %token is a token
 
+  (%anyAddr, %float, %token, %allocation) = begin_apply %0() : $@yield_once_2 () -> (@yields @inout %Any, @yields Float)
+  // %anyAddr : $*Any
+  // %float : $Float
+  // %token is a token
+  // %allocation is a pointer to a token
+
 Transfers control to coroutine ``%0``, passing it the given arguments.
 The rules for the application generally follow the rules for ``apply``,
 except:
 
-- the callee value must have a ``yield_once`` coroutine type,
+- the callee value must have be of single-yield coroutine type (``yield_once``
+  or ``yield_once_2``)
 
 - control returns to this function not when the coroutine performs a
   ``return``, but when it performs a ``yield``, and
 
 - the instruction results are derived from the yields of the coroutine
   instead of its normal results.
 
-The final result of a ``begin_apply`` is a "token", a special value which
-can only be used as the operand of an ``end_apply`` or ``abort_apply``
-instruction.  Before this second instruction is executed, the coroutine
-is said to be "suspended", and the token represents a reference to its
-suspended activation record.
+The final (in the case of ``@yield_once``) or penultimate (in the case of
+``@yield_once_2``) result of a ``begin_apply`` is a "token", a special value
+which can only be used as the operand of an ``end_apply`` or ``abort_apply``
+instruction.  Before this second instruction is executed, the coroutine is said
+to be "suspended", and the token represents a reference to its suspended
+activation record.
+
+If the coroutine's kind ``yield_once_2``, its final result is an address of a
+"token", representing the allocation done by the callee coroutine.  It can only
+be used as the operand of a ``dealloc_stack`` which must appear after the
+coroutine is resumed.
 
 The other results of the instruction correspond to the yields in the
 coroutine type.  In general, the rules of a yield are similar to the rules

include/swift/SIL/SILInstruction.h

@@ -847,6 +847,9 @@ class SILInstruction : public llvm::ilist_node<SILInstruction> {
   /// instruction.
   bool isAllocatingStack() const;
 
+  /// The stack allocation produced by the instruction, if any.
+  SILValue getStackAllocation() const;
+
   /// Returns true if this is the deallocation of a stack allocating instruction.
   /// The first operand must be the allocating instruction.
   bool isDeallocatingStack() const;
@@ -3256,12 +3259,20 @@ class BeginApplyInst final
   }
 
   MultipleValueInstructionResult *getTokenResult() const {
+    return const_cast<MultipleValueInstructionResult *>(
+        &getAllResultsBuffer().drop_back(isCalleeAllocated() ? 1 : 0).back());
+  }
+
+  MultipleValueInstructionResult *getCalleeAllocationResult() const {
+    if (!isCalleeAllocated()) {
+      return nullptr;
+    }
     return const_cast<MultipleValueInstructionResult *>(
              &getAllResultsBuffer().back());
   }
 
   SILInstructionResultArray getYieldedValues() const {
-    return getAllResultsBuffer().drop_back();
+    return getAllResultsBuffer().drop_back(isCalleeAllocated() ? 2 : 1);
   }
 
   void getCoroutineEndPoints(

lib/IRGen/GenCall.cpp

@@ -4635,6 +4635,16 @@ void irgen::emitDeallocYieldManyCoroutineBuffer(IRGenFunction &IGF,
   IGF.Builder.CreateLifetimeEnd(buffer, bufferSize);
 }
 
+void irgen::emitDeallocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
+                                                llvm::Value *allocation) {
+  if (IGF.IGM.IRGen.Opts.EmitYieldOnce2AsYieldOnce) {
+    assert(!allocation);
+    return;
+  }
+  assert(allocation);
+  llvm::report_fatal_error("unimplemented");
+}
+
 Address irgen::emitAllocAsyncContext(IRGenFunction &IGF,
                                      llvm::Value *sizeValue) {
   auto alignment = IGF.IGM.getAsyncContextAlignment();

lib/IRGen/GenCall.h

@@ -203,6 +203,8 @@ namespace irgen {
 
   Address emitAllocYieldOnceCoroutineBuffer(IRGenFunction &IGF);
   void emitDeallocYieldOnceCoroutineBuffer(IRGenFunction &IGF, Address buffer);
+  void emitDeallocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
+                                           llvm::Value *allocation);
   void
   emitYieldOnceCoroutineEntry(IRGenFunction &IGF,
                               CanSILFunctionType coroutineType,

lib/IRGen/IRGenSIL.cpp

@@ -366,6 +366,7 @@ class IRGenSILFunction :
   llvm::DenseMap<SILValue, LoweredValue> LoweredValues;
   llvm::DenseMap<SILValue, StackAddress> LoweredPartialApplyAllocations;
   llvm::DenseMap<SILType, LoweredValue> LoweredUndefs;
+  llvm::DenseMap<SILValue, llvm::Value *> LoweredSingleYieldFrames;
 
   /// All alloc_ref instructions which allocate the object on the stack.
   llvm::SmallPtrSet<SILInstruction *, 8> StackAllocs;
@@ -6264,6 +6265,13 @@ void IRGenSILFunction::visitDeallocStackInst(swift::DeallocStackInst *i) {
     emitDeallocateDynamicAlloca(stackAddr);
     return;
   }
+  if (isaResultOf<BeginApplyInst>(i->getOperand())) {
+    auto *mvi = getAsResultOf<BeginApplyInst>(i->getOperand());
+    auto *bai = cast<BeginApplyInst>(mvi->getParent());
+    emitDeallocYieldOnce2CoroutineFrame(
+        *this, LoweredSingleYieldFrames[bai->getCalleeAllocationResult()]);
+    return;
+  }
 
   auto allocatedType = i->getOperand()->getType();
   const TypeInfo &allocatedTI = getTypeInfo(allocatedType);

lib/SIL/IR/SILInstruction.cpp

@@ -1283,6 +1283,10 @@ bool SILInstruction::isAllocatingStack() const {
       && !PA->getFunction()->hasOwnership();
   }
 
+  if (auto *BAI = dyn_cast<BeginApplyInst>(this)) {
+    return BAI->isCalleeAllocated();
+  }
+
   if (auto *BI = dyn_cast<BuiltinInst>(this)) {
     if (BI->getBuiltinKind() == BuiltinValueKind::StackAlloc ||
         BI->getBuiltinKind() == BuiltinValueKind::UnprotectedStackAlloc) {
@@ -1293,6 +1297,17 @@ bool SILInstruction::isAllocatingStack() const {
   return false;
 }
 
+SILValue SILInstruction::getStackAllocation() const {
+  if (!isAllocatingStack()) {
+    return {};
+  }
+
+  if (auto *bai = dyn_cast<BeginApplyInst>(this)) {
+    return bai->getCalleeAllocationResult();
+  }
+  return cast<SingleValueInstruction>(this);
+}
+
 bool SILInstruction::isDeallocatingStack() const {
   if (isa<DeallocStackInst>(this) ||
       isa<DeallocStackRefInst>(this) ||

lib/SIL/IR/SILInstructions.cpp

@@ -665,14 +665,19 @@ BeginApplyInst *BeginApplyInst::create(
             : SILModuleConventions(parentFunction.getModule())));
   }
 
-  resultTypes.push_back(
-      SILType::getSILTokenType(parentFunction.getASTContext()));
+  auto tokenTy = SILType::getSILTokenType(parentFunction.getASTContext());
+  resultTypes.push_back(tokenTy);
   // The begin_apply token represents the borrow scope of all owned and
   // guaranteed call arguments. Although SILToken is (currently) trivially
   // typed, it must have guaranteed ownership so end_apply and abort_apply will
   // be recognized as lifetime-ending uses.
   resultOwnerships.push_back(OwnershipKind::Guaranteed);
 
+  if (substCalleeType->isCalleeAllocatedCoroutine()) {
+    resultTypes.push_back(tokenTy.getAddressType());
+    resultOwnerships.push_back(OwnershipKind::None);
+  }
+
   SmallVector<SILValue, 32> typeDependentOperands;
   collectTypeDependentOperands(typeDependentOperands, parentFunction,
                                substCalleeType, subs);

lib/SIL/Verifier/SILVerifier.cpp

@@ -1688,8 +1688,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
   }
 
   static bool isLegalSILTokenProducer(SILValue value) {
-    if (auto *baResult = isaResultOf<BeginApplyInst>(value))
-      return baResult->isBeginApplyToken();
+    if (auto *baResult = isaResultOf<BeginApplyInst>(value)) {
+      auto *bai = cast<BeginApplyInst>(baResult->getParent());
+      return value == bai->getTokenResult() ||
+             value == bai->getCalleeAllocationResult();
+    }
 
     if (isa<OpenPackElementInst>(value))
       return true;
@@ -3772,12 +3775,23 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
   }
 
   void checkDeallocStackInst(DeallocStackInst *DI) {
+    auto isTokenFromCalleeAllocatedBeginApply = [](SILValue value) -> bool {
+      auto *inst = value->getDefiningInstruction();
+      if (!inst)
+        return false;
+      auto *bai = dyn_cast<BeginApplyInst>(inst);
+      if (!bai)
+        return false;
+      return value == bai->getCalleeAllocationResult();
+    };
     require(isa<SILUndef>(DI->getOperand()) ||
                 isa<AllocStackInst>(DI->getOperand()) ||
                 isa<AllocVectorInst>(DI->getOperand()) ||
                 (isa<PartialApplyInst>(DI->getOperand()) &&
-                 cast<PartialApplyInst>(DI->getOperand())->isOnStack()),
-            "Operand of dealloc_stack must be an alloc_stack, alloc_vector or partial_apply "
+                 cast<PartialApplyInst>(DI->getOperand())->isOnStack()) ||
+                (isTokenFromCalleeAllocatedBeginApply(DI->getOperand())),
+            "Operand of dealloc_stack must be an alloc_stack, alloc_vector or "
+            "partial_apply "
             "[stack]");
   }
   void checkDeallocPackInst(DeallocPackInst *DI) {
@@ -6768,7 +6782,7 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     };
 
     struct BBState {
-      std::vector<SingleValueInstruction*> Stack;
+      std::vector<SILValue> Stack;
 
       /// Contents: BeginAccessInst*, BeginApplyInst*.
       std::set<SILInstruction*> ActiveOps;
@@ -6818,9 +6832,15 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
         }
 
         if (i.isAllocatingStack()) {
-          state.Stack.push_back(cast<SingleValueInstruction>(&i));
-
-        } else if (i.isDeallocatingStack()) {
+          if (auto *BAI = dyn_cast<BeginApplyInst>(&i)) {
+            state.Stack.push_back(BAI->getCalleeAllocationResult());
+          } else {
+            state.Stack.push_back(cast<SingleValueInstruction>(&i));
+          }
+          // Not "else if": begin_apply both allocates stack and begins an
+          // operation.
+        }
+        if (i.isDeallocatingStack()) {
           SILValue op = i.getOperand(0);
           while (auto *mvi = dyn_cast<MoveValueInst>(op)) {
             op = mvi->getOperand();

lib/SILGen/SILGenApply.cpp

@@ -5056,6 +5056,10 @@ SILGenFunction::emitBeginApply(SILLocation loc, ManagedValue fn,
                /*indirect results*/ {}, /*indirect errors*/ {},
                rawResults, ExecutorBreadcrumb());
 
+  if (substFnType->isCalleeAllocatedCoroutine()) {
+    auto allocation = rawResults.pop_back_val();
+    enterDeallocStackCleanup(allocation);
+  }
   auto token = rawResults.pop_back_val();
   auto yieldValues = llvm::ArrayRef(rawResults);
 

lib/SILOptimizer/Utils/LoopUtils.cpp

@@ -231,8 +231,8 @@ bool swift::canDuplicateLoopInstruction(SILLoop *L, SILInstruction *I) {
 
   // The deallocation of a stack allocation must be in the loop, otherwise the
   // deallocation will be fed by a phi node of two allocations.
-  if (I->isAllocatingStack()) {
-    for (auto *UI : cast<SingleValueInstruction>(I)->getUses()) {
+  if (auto allocation = I->getStackAllocation()) {
+    for (auto *UI : allocation->getUses()) {
       if (UI->getUser()->isDeallocatingStack()) {
         if (!L->contains(UI->getUser()->getParent()))
           return false;
@@ -246,6 +246,8 @@ bool swift::canDuplicateLoopInstruction(SILLoop *L, SILInstruction *I) {
       SILValue address = dealloc->getOperand();
       if (isa<AllocStackInst>(address) || isa<PartialApplyInst>(address))
         alloc = cast<SingleValueInstruction>(address);
+      else if (isaResultOf<BeginApplyInst>(address))
+        alloc = cast<MultipleValueInstructionResult>(address)->getParent();
     }
     if (auto *dealloc = dyn_cast<DeallocStackRefInst>(I))
       alloc = dealloc->getAllocRef();