Merge pull request #76942 from nate-chandler/general-coro/20240927/1

[CoroutineAccessors] SIL represents callee alloc.

Author: nate-chandler

docs/ABI/Mangling.rst

@@ -856,6 +856,7 @@ mangled in to disambiguate.
   FUNC-REPRESENTATION ::= 'W'                // protocol witness
 
   COROUTINE-KIND ::= 'A'                     // yield-once coroutine
+  COROUTINE-KIND ::= 'I'                     // yield-once-2 coroutine
   COROUTINE-KIND ::= 'G'                     // yield-many coroutine
 
   #if SWIFT_RUNTIME_VERSION >= 5.5

docs/SIL.rst

@@ -789,13 +789,16 @@ coroutine, where a normal function return would need to transfer ownership of
 its return value, since a normal function's context ceases to exist and be able
 to maintain ownership of the value after it returns.
 
-To support these concepts, SIL supports two kinds of coroutine:
-``@yield_many`` and ``@yield_once``. Either of these attributes may be
+To support these concepts, SIL supports two flavors: single-yield and
+multi-yield.  These two flavors correspond to three kinds.  A multi-yield
+coroutine is of kind ``@yield_many``. A single-yield coroutine is of kind
+either ``@yield_once`` or ``@yield_once_2``. Any of these attributes may be
 written before a function type to indicate that it is a coroutine type.
-``@yield_many`` and ``@yield_once`` coroutines are allowed to also be
-``@async``. (Note that ``@async`` functions are not themselves modeled
-explicitly as coroutines in SIL, although the implementation may use a coroutine
-lowering strategy.)
+
+Both single-yield and multi-yield coroutines are allowed to also be ``@async``.
+(Note that ``@async`` functions are not themselves modeled explicitly as
+coroutines in SIL, although the implementation may use a coroutine lowering
+strategy.)
 
 A coroutine type may declare any number of *yielded values*, which is to
 say, values which are provided to the caller at a yield point.  Yielded
@@ -816,9 +819,10 @@ calling a yield-many coroutine of any kind.
 Coroutines may contain the special ``yield`` and ``unwind``
 instructions.
 
-A ``@yield_many`` coroutine may yield as many times as it desires.
-A ``@yield_once`` coroutine may yield exactly once before returning,
-although it may also ``throw`` before reaching that point.
+A multi-yield (``@yield_many``) coroutine may yield as many times as it desires.
+A single-yield (``@yield_once`` or ``@yield_once_2``) coroutine may yield
+exactly once before returning, although it may also ``throw`` before reaching
+that point.
 
 Variadic Generics
 `````````````````
@@ -6128,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
@@ -8376,9 +8393,10 @@ the current function was invoked with a ``try_apply`` instruction, control
 resumes at the normal destination, and the value of the basic block argument
 will be the operand of this ``return`` instruction.
 
-If the current function is a ``yield_once`` coroutine, there must not be
-a path from the entry block to a ``return`` which does not pass through
-a ``yield`` instruction. This rule does not apply in the ``raw`` SIL stage.
+If the current function is a single-yield coroutine (``yield_once`` or
+``yield_once_2``), there must not be a path from the entry block to a
+``return`` which does not pass through a ``yield`` instruction. This rule does
+not apply in the ``raw`` SIL stage.
 
 ``return`` does not retain or release its operand or any other values.
 
@@ -8446,9 +8464,10 @@ The ``resume`` and ``unwind`` destination blocks must be uniquely
 referenced by the ``yield`` instruction.  This prevents them from becoming
 critical edges.
 
-In a ``yield_once`` coroutine, there must not be a control flow path leading
-from the ``resume`` edge to another ``yield`` instruction in this function.
-This rule does not apply in the ``raw`` SIL stage.
+In a single-yield coroutine (``yield_once`` or ``yield_once_2``), there must
+not be a control flow path leading from the ``resume`` edge to another
+``yield`` instruction in this function. This rule does not apply in the ``raw``
+SIL stage.
 
 There must not be a control flow path leading from the ``unwind`` edge to
 a ``return`` instruction, to a ``throw`` instruction, or to any block

include/swift/AST/IRGenOptions.h

@@ -165,6 +165,9 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
   /// Resumption functions from yield-once coroutines.
   PointerAuthSchema YieldOnceResumeFunctions;
 
+  /// Resumption functions from yield-once-2 coroutines.
+  PointerAuthSchema YieldOnce2ResumeFunctions;
+
   /// Resumption functions from yield-many coroutines.
   PointerAuthSchema YieldManyResumeFunctions;
 
@@ -482,6 +485,9 @@ class IRGenOptions {
 
   unsigned EmitAsyncFramePushPopMetadata : 1;
 
+  // Whether to use the yield_once ABI when emitting yield_once_2 coroutines.
+  unsigned EmitYieldOnce2AsYieldOnce : 1;
+
   // Whether to force emission of a frame for all async functions
   // (LLVM's 'frame-pointer=all').
   unsigned AsyncFramePointerAll : 1;
@@ -574,10 +580,10 @@ class IRGenOptions {
         EmitGenericRODatas(true), NoPreallocatedInstantiationCaches(false),
         DisableReadonlyStaticObjects(false), CollocatedMetadataFunctions(false),
         ColocateTypeDescriptors(true), UseRelativeProtocolWitnessTables(false),
-        UseFragileResilientProtocolWitnesses(false),
-        EnableHotColdSplit(false), EmitAsyncFramePushPopMetadata(false),
-        AsyncFramePointerAll(false),
-        CmdArgs(), SanitizeCoverage(llvm::SanitizerCoverageOptions()),
+        UseFragileResilientProtocolWitnesses(false), EnableHotColdSplit(false),
+        EmitAsyncFramePushPopMetadata(false), EmitYieldOnce2AsYieldOnce(true),
+        AsyncFramePointerAll(false), CmdArgs(),
+        SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All),
         PlatformCCallingConvention(llvm::CallingConv::C), UseCASBackend(false),
         CASObjMode(llvm::CASBackendMode::Native) {

include/swift/AST/TypeAttr.def

@@ -101,6 +101,7 @@ SIMPLE_SIL_TYPE_ATTR(pseudogeneric, Pseudogeneric)
 SIMPLE_SIL_TYPE_ATTR(unimplementable, Unimplementable)
 SIMPLE_SIL_TYPE_ATTR(yields, Yields)
 SIMPLE_SIL_TYPE_ATTR(yield_once, YieldOnce)
+SIMPLE_SIL_TYPE_ATTR(yield_once_2, YieldOnce2)
 SIMPLE_SIL_TYPE_ATTR(yield_many, YieldMany)
 SIMPLE_SIL_TYPE_ATTR(captures_generics, CapturesGenerics)
 // Used at the SIL level to mark a type as moveOnly.

include/swift/AST/Types.h

@@ -4823,11 +4823,18 @@ enum class SILCoroutineKind : uint8_t {
   /// It must not have normal results and may have arbitrary yield results.
   YieldOnce,
 
+  /// This function is a yield-once coroutine (used by read and modify
+  /// accessors).  It has the following differences from YieldOnce:
+  /// - it does not observe errors thrown by its caller (unless the feature
+  /// CoroutineAccessorsUnwindOnCallerError is enabled)
+  /// - it uses the callee-allocated ABI
+  YieldOnce2,
+
   /// This function is a yield-many coroutine (used by e.g. generators).
   /// It must not have normal results and may have arbitrary yield results.
   YieldMany,
 };
-  
+
 class SILFunctionConventions;
 
 
@@ -5039,6 +5046,17 @@ class SILFunctionType final
   SILCoroutineKind getCoroutineKind() const {
     return SILCoroutineKind(Bits.SILFunctionType.CoroutineKind);
   }
+  /// Whether this coroutine's ABI is callee-allocated.
+  bool isCalleeAllocatedCoroutine() const {
+    switch (getCoroutineKind()) {
+    case SILCoroutineKind::None:
+    case SILCoroutineKind::YieldOnce:
+    case SILCoroutineKind::YieldMany:
+      return false;
+    case SILCoroutineKind::YieldOnce2:
+      return true;
+    }
+  }
 
   bool isSendable() const { return getExtInfo().isSendable(); }
   bool isUnimplementable() const { return getExtInfo().isUnimplementable(); }

include/swift/Basic/Features.def

@@ -427,6 +427,12 @@ EXPERIMENTAL_FEATURE(UnspecifiedMeansMainActorIsolated, false)
 /// modify/read single-yield coroutines
 SUPPRESSIBLE_EXPERIMENTAL_FEATURE(CoroutineAccessors, true)
 
+/// modify/read single-yield coroutines always execute code post-yield code
+EXPERIMENTAL_FEATURE(CoroutineAccessorsUnwindOnCallerError, false)
+
+/// modify/read coroutines use the callee-allocated ABI
+EXPERIMENTAL_FEATURE(CoroutineAccessorsAllocateInCallee, false)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE

include/swift/Demangling/TypeDecoder.h

@@ -51,6 +51,7 @@ enum class ImplMetatypeRepresentation {
 enum class ImplCoroutineKind {
   None,
   YieldOnce,
+  YieldOnce2,
   YieldMany,
 };
 
@@ -1097,6 +1098,8 @@ class TypeDecoder {
             return MAKE_NODE_TYPE_ERROR0(child, "expected text");
           if (child->getText() == "yield_once") {
             coroutineKind = ImplCoroutineKind::YieldOnce;
+          } else if (child->getText() == "yield_once_2") {
+            coroutineKind = ImplCoroutineKind::YieldOnce2;
           } else if (child->getText() == "yield_many") {
             coroutineKind = ImplCoroutineKind::YieldMany;
           } else

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;
@@ -3251,13 +3254,25 @@ class BeginApplyInst final
 public:
   using MultipleValueInstructionTrailingObjects::totalSizeToAlloc;
 
+  bool isCalleeAllocated() const {
+    return getSubstCalleeType()->isCalleeAllocatedCoroutine();
+  }
+
   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/AST/ASTDemangler.cpp

@@ -582,6 +582,8 @@ getCoroutineKind(ImplCoroutineKind kind) {
     return SILCoroutineKind::None;
   case ImplCoroutineKind::YieldOnce:
     return SILCoroutineKind::YieldOnce;
+  case ImplCoroutineKind::YieldOnce2:
+    return SILCoroutineKind::YieldOnce2;
   case ImplCoroutineKind::YieldMany:
     return SILCoroutineKind::YieldMany;
   }

lib/AST/ASTMangler.cpp

@@ -2149,6 +2149,9 @@ void ASTMangler::appendImplFunctionType(SILFunctionType *fn,
   case SILCoroutineKind::YieldOnce:
     OpArgs.push_back('A');
     break;
+  case SILCoroutineKind::YieldOnce2:
+    OpArgs.push_back('I');
+    break;
   case SILCoroutineKind::YieldMany:
     OpArgs.push_back('G');
     break;