Merge pull request #37938 from jckarter/async-let-multi-suspend

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

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

@@ -793,18 +793,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/AST/DiagnosticsSema.def

@@ -3711,6 +3711,10 @@ WARNING(partial_application_of_function_invalid_swift4,none,
       "be an error in future Swift versions",
       (unsigned))
 
+// Cannot capture `async let`
+ERROR(capture_async_let_not_supported,none,
+      "capturing 'async let' variables is not supported", ())
+
 ERROR(self_assignment_var,none,
       "assigning a variable to itself", ())
 ERROR(self_assignment_prop,none,

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,121 @@ 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.
 ///

include/swift/Runtime/RuntimeFunctions.def

@@ -1700,11 +1700,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

@@ -2812,9 +2812,11 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     return getDistributedActorInitDestroy(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: