Handle multiple awaits and suspend-on-exit for async let tasks.

Change the code generation patterns for `async let` bindings to use an ABI based on the following
functions:

- `swift_asyncLet_begin`, which starts an `async let` child task, but which additionally
  now associates the `async let` with a caller-owned buffer to receive the result of the task.
  This is intended to allow the task to emplace its result in caller-owned memory, allowing the
  child task to be deallocated after completion without invalidating the result buffer.
- `swift_asyncLet_get[_throwing]`, which replaces `swift_asyncLet_wait[_throwing]`. Instead of
  returning a copy of the value, this entry point concerns itself with populating the local buffer.
  If the buffer hasn't been populated, then it awaits completion of the task and emplaces the
  result in the buffer; otherwise, it simply returns. The caller can then read the result out of
  its owned memory. These entry points are intended to be used before every read from the
  `async let` binding, after which point the local buffer is guaranteed to contain an initialized
  value.
- `swift_asyncLet_finish`, which replaces `swift_asyncLet_end`. Unlike `_end`, this variant
  is async and will suspend the parent task after cancelling the child to ensure it finishes
  before cleaning up. The local buffer will also be deinitialized if necessary. This is intended
  to be used on exit from an `async let` scope, to handle cleaning up the local buffer if necessary
  as well as cancelling, awaiting, and deallocating the child task.
- `swift_asyncLet_consume[_throwing]`, which combines `get` and `finish`. This will await completion
  of the task, leaving the result value in the result buffer (or propagating the error, if it
  throws), while destroying and deallocating the child task. This is intended as an optimization
  for reading `async let` variables that are read exactly once by their parent task.

To avoid an epoch break with existing swiftinterfaces and ABI clients, the old builtins and entry
points are kept intact for now, but SILGen now only generates code using the new interface.

This new interface fixes several issues with the old async let codegen, including use-after-free
crashes if the `async let` was never awaited, and the inability to read from an `async let` variable
more than once.

rdar://77855176

Author: jckarter

include/swift/ABI/AsyncLet.h

@@ -36,15 +36,29 @@ class alignas(Alignment_AsyncLet) AsyncLet {
   constexpr AsyncLet()
     : PrivateData{} {}
 
-  // FIXME: not sure how many words we should reserve
   void *PrivateData[NumWords_AsyncLet];
 
   // TODO: we could offer a "was awaited on" check here
 
   /// Returns the child task that is associated with this async let.
   /// The tasks completion is used to fulfil the value represented by this async let.
   AsyncTask *getTask() const;
-
+  
+  // The compiler preallocates a large fixed space for the `async let`, with the
+  // intent that most of it be used for the child task context. The next two
+  // methods return the address and size of that space.
+  
+  /// Return a pointer to the unused space within the async let block.
+  void *getPreallocatedSpace();
+  
+  /// Return the size of the unused space within the async let block.
+  static size_t getSizeOfPreallocatedSpace();
+  
+  /// Was the task allocated out of the parent's allocator?
+  bool didAllocateFromParentTask();
+  
+  /// Flag that the task was allocated from the parent's allocator.
+  void setDidAllocateFromParentTask(bool value = true);
 };
 
 } // end namespace swift

include/swift/ABI/MetadataValues.h

@@ -53,8 +53,8 @@ enum {
   /// The number of words in a task group.
   NumWords_TaskGroup = 32,
 
-  /// The number of words in an AsyncLet (flags + task pointer)
-  NumWords_AsyncLet = 8, // TODO: not sure how much is enough, these likely could be pretty small
+  /// The number of words in an AsyncLet (flags + child task context & allocation)
+  NumWords_AsyncLet = 80, // 640 bytes ought to be enough for anyone
 };
 
 struct InProcess;
@@ -2145,8 +2145,11 @@ enum class TaskOptionRecordKind : uint8_t {
   Executor  = 0,
   /// Request a child task to be part of a specific task group.
   TaskGroup = 1,
+  /// DEPRECATED. AsyncLetWithBuffer is used instead.
   /// Request a child task for an 'async let'.
   AsyncLet  = 2,
+  /// Request a child task for an 'async let'.
+  AsyncLetWithBuffer = 3,
 };
 
 /// Flags for cancellation records.

include/swift/ABI/TaskOptions.h

@@ -98,11 +98,12 @@ class ExecutorTaskOptionRecord : public TaskOptionRecord {
   }
 };
 
+/// DEPRECATED. AsyncLetWithBufferTaskOptionRecord is used instead.
 /// Task option to specify that the created task is for an 'async let'.
 class AsyncLetTaskOptionRecord : public TaskOptionRecord {
   AsyncLet *asyncLet;
 
-  public:
+public:
   AsyncLetTaskOptionRecord(AsyncLet *asyncLet)
     : TaskOptionRecord(TaskOptionRecordKind::AsyncLet),
       asyncLet(asyncLet) {}
@@ -116,6 +117,30 @@ class AsyncLetTaskOptionRecord : public TaskOptionRecord {
   }
 };
 
+class AsyncLetWithBufferTaskOptionRecord : public TaskOptionRecord {
+  AsyncLet *asyncLet;
+  void *resultBuffer;
+
+public:
+  AsyncLetWithBufferTaskOptionRecord(AsyncLet *asyncLet,
+                                     void *resultBuffer)
+    : TaskOptionRecord(TaskOptionRecordKind::AsyncLetWithBuffer),
+      asyncLet(asyncLet),
+      resultBuffer(resultBuffer) {}
+
+  AsyncLet *getAsyncLet() const {
+    return asyncLet;
+  }
+  
+  void *getResultBuffer() const {
+    return resultBuffer;
+  }
+
+  static bool classof(const TaskOptionRecord *record) {
+    return record->getKind() == TaskOptionRecordKind::AsyncLetWithBuffer;
+  }
+};
+
 } // end namespace swift
 
 #endif

include/swift/AST/Builtins.def

@@ -782,18 +782,43 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(CancelAsyncTask, "cancelAsyncTask", "", Speci
 ///     __owned @Sendable @escaping () async throws -> T
 /// ) -> Builtin.RawPointer
 ///
+/// DEPRECATED. startAsyncLetWithLocalBuffer is used instead.
+///
 /// Create, initialize and start a new async-let and associated task.
 /// Returns an AsyncLet* that must be passed to endAsyncLet for destruction.
 BUILTIN_MISC_OPERATION(StartAsyncLet, "startAsyncLet", "", Special)
 
-/// asyncLetEnd(): (Builtin.RawPointer) -> Void
+/// startAsyncLetWithLocalBuffer()<T>: (
+///     __owned @Sendable @escaping () async throws -> T,
+///     _ resultBuf: Builtin.RawPointer
+/// ) -> Builtin.RawPointer
+///
+/// Create, initialize and start a new async-let and associated task, with a
+/// locally-allocated buffer assigned to receive the result if the task
+/// completes.
+/// Returns an AsyncLet* that must be passed to endAsyncLetLifetime for
+/// destruction.
+BUILTIN_MISC_OPERATION(StartAsyncLetWithLocalBuffer, "startAsyncLetWithLocalBuffer", "", Special)
+
+/// endAsyncLet(): (Builtin.RawPointer) -> Void
+///
+/// DEPRECATED. The swift_asyncLet_finish intrinsic and endAsyncLetLifetime
+/// builtin are used instead.
 ///
 /// Ends and destroys an async-let.
 /// The ClosureLifetimeFixup pass adds a second operand to the builtin to
 /// ensure that optimizations keep the stack-allocated closure arguments alive
 /// until the endAsyncLet.
 BUILTIN_MISC_OPERATION_WITH_SILGEN(EndAsyncLet, "endAsyncLet", "", Special)
 
+/// endAsyncLetLifetime(): (Builtin.RawPointer) -> Void
+///
+/// Marks the end of an async-let's lifetime.
+/// The ClosureLifetimeFixup pass adds a second operand to the builtin to
+/// ensure that optimizations keep the stack-allocated closure arguments alive
+/// until the endAsyncLet.
+BUILTIN_MISC_OPERATION(EndAsyncLetLifetime, "endAsyncLetLifetime", "", Special)
+
 /// createAsyncTask(): (
 ///     Int,                            // task-creation flags
 ///     @escaping () async throws -> T  // function

include/swift/Runtime/Concurrency.h

@@ -240,6 +240,7 @@ bool swift_taskGroup_isCancelled(TaskGroup *group);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_taskGroup_isEmpty(TaskGroup *group);
 
+/// DEPRECATED. swift_asyncLet_begin is used instead.
 /// Its Swift signature is
 ///
 /// \code
@@ -255,12 +256,31 @@ void swift_asyncLet_start(AsyncLet *alet,
                           const Metadata *futureResultType,
                           void *closureEntryPoint, HeapObject *closureContext);
 
+/// Begin an async let child task.
+/// Its Swift signature is
+///
+/// \code
+/// func swift_asyncLet_start<T>(
+///     asyncLet: Builtin.RawPointer,
+///     options: Builtin.RawPointer?,
+///     operation: __owned @Sendable () async throws -> T,
+///     resultBuffer: Builtin.RawPointer
+/// )
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_asyncLet_begin(AsyncLet *alet,
+                          TaskOptionRecord *options,
+                          const Metadata *futureResultType,
+                          void *closureEntryPoint, HeapObject *closureContext,
+                          void *resultBuffer);
+
 /// This matches the ABI of a closure `<T>(Builtin.RawPointer) async -> T`
 using AsyncLetWaitSignature =
     SWIFT_CC(swiftasync)
     void(OpaqueValue *,
          SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
 
+/// DEPRECATED. swift_asyncLet_get is used instead.
 /// Wait for a non-throwing async-let to complete.
 ///
 /// This can be called from any thread. Its Swift signature is
@@ -276,6 +296,7 @@ void swift_asyncLet_wait(OpaqueValue *,
                          AsyncLet *, TaskContinuationFunction *,
                          AsyncContext *);
 
+/// DEPRECATED. swift_asyncLet_get_throwing is used instead.
 /// Wait for a potentially-throwing async-let to complete.
 ///
 /// This can be called from any thread. Its Swift signature is
@@ -292,6 +313,7 @@ void swift_asyncLet_wait_throwing(OpaqueValue *,
                                   ThrowingTaskFutureWaitContinuationFunction *,
                                   AsyncContext *);
 
+/// DEPRECATED. swift_asyncLet_finish is used instead.
 /// Its Swift signature is
 ///
 /// \code
@@ -300,6 +322,134 @@ void swift_asyncLet_wait_throwing(OpaqueValue *,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_asyncLet_end(AsyncLet *alet);
 
+/// Get the value of a non-throwing async-let, awaiting the result if necessary.
+///
+/// This can be called from any thread. Its Swift signature is
+///
+/// \code
+/// func swift_asyncLet_get(
+///     _ asyncLet: Builtin.RawPointer,
+///     _ resultBuffer: Builtin.RawPointer
+/// ) async
+/// \endcode
+///
+/// \c result points at the variable storage for the binding. It is
+/// uninitialized until the first call to \c swift_asyncLet_get or
+/// \c swift_asyncLet_get_throwing. That first call initializes the storage
+/// with the result of the child task. Subsequent calls do nothing and leave
+/// the value in place.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
+void swift_asyncLet_get(SWIFT_ASYNC_CONTEXT AsyncContext *,
+                        AsyncLet *,
+                        void *,
+                        TaskContinuationFunction *,
+                        AsyncContext *);
+
+/// Get the value of a throwing async-let, awaiting the result if necessary.
+///
+/// This can be called from any thread. Its Swift signature is
+///
+/// \code
+/// func swift_asyncLet_get_throwing(
+///     _ asyncLet: Builtin.RawPointer,
+///     _ resultBuffer: Builtin.RawPointer
+/// ) async throws
+/// \endcode
+///
+/// \c result points at the variable storage for the binding. It is
+/// uninitialized until the first call to \c swift_asyncLet_get or
+/// \c swift_asyncLet_get_throwing. That first call initializes the storage
+/// with the result of the child task. Subsequent calls do nothing and leave
+/// the value in place. A pointer to the storage inside the child task is
+/// returned if the task completes successfully, otherwise the error from the
+/// child task is thrown.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
+void swift_asyncLet_get_throwing(SWIFT_ASYNC_CONTEXT AsyncContext *,
+                                 AsyncLet *,
+                                 void *,
+                                 ThrowingTaskFutureWaitContinuationFunction *,
+                                 AsyncContext *);
+
+/// Exit the scope of an async-let binding. If the task is still running, it
+/// is cancelled, and we await its completion; otherwise, we destroy the
+/// value in the variable storage.
+///
+/// Its Swift signature is
+///
+/// \code
+/// func swift_asyncLet_finish(_ asyncLet: Builtin.RawPointer,
+///                            _ resultBuffer: Builtin.RawPointer) async
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
+void swift_asyncLet_finish(SWIFT_ASYNC_CONTEXT AsyncContext *,
+                           AsyncLet *,
+                           void *,
+                           TaskContinuationFunction *,
+                           AsyncContext *);
+
+/// Get the value of a non-throwing async-let, awaiting the result if necessary,
+/// and then destroy the child task. The result buffer is left initialized after
+/// returning.
+///
+/// This can be called from any thread. Its Swift signature is
+///
+/// \code
+/// func swift_asyncLet_get(
+///     _ asyncLet: Builtin.RawPointer,
+///     _ resultBuffer: Builtin.RawPointer
+/// ) async
+/// \endcode
+///
+/// \c result points at the variable storage for the binding. It is
+/// uninitialized until the first call to \c swift_asyncLet_get or
+/// \c swift_asyncLet_get_throwing. The child task will be invalidated after
+/// this call, so the `async let` can not be gotten or finished afterward.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
+void swift_asyncLet_consume(SWIFT_ASYNC_CONTEXT AsyncContext *,
+                            AsyncLet *,
+                            void *,
+                            TaskContinuationFunction *,
+                            AsyncContext *);
+
+/// Get the value of a throwing async-let, awaiting the result if necessary,
+/// and then destroy the child task. The result buffer is left initialized after
+/// returning.
+///
+/// This can be called from any thread. Its Swift signature is
+///
+/// \code
+/// func swift_asyncLet_get_throwing(
+///     _ asyncLet: Builtin.RawPointer,
+///     _ resultBuffer: Builtin.RawPointer
+/// ) async throws
+/// \endcode
+///
+/// \c result points at the variable storage for the binding. It is
+/// uninitialized until the first call to \c swift_asyncLet_get or
+/// \c swift_asyncLet_get_throwing. That first call initializes the storage
+/// with the result of the child task. Subsequent calls do nothing and leave
+/// the value in place. The child task will be invalidated after
+/// this call, so the `async let` can not be gotten or finished afterward.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
+void swift_asyncLet_consume_throwing(SWIFT_ASYNC_CONTEXT AsyncContext *,
+                                     AsyncLet *,
+                                     void *,
+                                     ThrowingTaskFutureWaitContinuationFunction *,
+                                     AsyncContext *);
+
+/// Returns true if the currently executing AsyncTask has a
+/// 'TaskGroupTaskStatusRecord' present.
+///
+/// This can be called from any thread.
+///
+/// Its Swift signature is
+///
+/// \code
+/// func swift_taskGroup_hasTaskGroupRecord()
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+bool swift_taskGroup_hasTaskGroupRecord();
+
 /// Add a status record to a task.  The record should not be
 /// modified while it is registered with a task.
 ///

include/swift/Runtime/RuntimeFunctions.def

@@ -1674,11 +1674,32 @@ FUNCTION(AsyncLetStart,
          swift_asyncLet_start, SwiftCC,
          ConcurrencyAvailability,
          RETURNS(VoidTy),
+         ARGS(SwiftAsyncLetPtrTy,         // AsyncLet*
+              SwiftTaskOptionRecordPtrTy, // options
+              TypeMetadataPtrTy,          // futureResultType
+              Int8PtrTy,                  // closureEntry
+              OpaquePtrTy                 // closureContext
+         ),
+         ATTRS(NoUnwind, ArgMemOnly))
+
+/// void swift_asyncLet_begin(
+///     AsyncLet *alet,
+///     TaskOptionRecord *options,
+///     const Metadata *futureResultType,
+///     void *closureEntryPoint,
+///     HeapObject *closureContext,
+///     void *resultBuffer
+/// );
+FUNCTION(AsyncLetBegin,
+         swift_asyncLet_begin, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
          ARGS(SwiftAsyncLetPtrTy,         // AsyncLet*
               SwiftTaskOptionRecordPtrTy, // options
               TypeMetadataPtrTy,          // futureResultType
               Int8PtrTy,                  // closureEntry
               OpaquePtrTy,                // closureContext
+              Int8PtrTy
          ),
          ATTRS(NoUnwind, ArgMemOnly))
 

lib/AST/Builtins.cpp

@@ -2790,9 +2790,11 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     return getDefaultActorInitDestroy(Context, Id);
 
   case BuiltinValueKind::StartAsyncLet:
+  case BuiltinValueKind::StartAsyncLetWithLocalBuffer:
     return getStartAsyncLet(Context, Id);
 
   case BuiltinValueKind::EndAsyncLet:
+  case BuiltinValueKind::EndAsyncLetLifetime:
     return getEndAsyncLet(Context, Id);
 
   case BuiltinValueKind::CreateTaskGroup: