[WIP][Concurrency] introduce options to wait operations

include/swift/ABI/TaskOptions.h

@@ -80,14 +80,14 @@ class TaskGroupTaskOptionRecord : public TaskOptionRecord {
 /// executor should be used instead, most often this may mean the global
 /// concurrent executor, or the enclosing actor's executor.
 class ExecutorTaskOptionRecord : public TaskOptionRecord {
-  ExecutorRef *Executor;
+  ExecutorRef Executor;
 
 public:
-  ExecutorTaskOptionRecord(ExecutorRef *executor)
+  ExecutorTaskOptionRecord(ExecutorRef executor)
     : TaskOptionRecord(TaskOptionRecordKind::Executor),
       Executor(executor) {}
 
-  ExecutorRef *getExecutor() const {
+  ExecutorRef getExecutor() const {
     return Executor;
   }
 };

include/swift/Runtime/Concurrency.h

@@ -130,12 +130,16 @@ swift_task_escalate(AsyncTask *task, JobPriority newPriority);
 /// This can be called from any thread. Its Swift signature is
 ///
 /// \code
-/// func swift_task_future_wait(on task: _owned Builtin.NativeObject) async
-///     -> Success
+/// func swift_task_future_wait(
+///     on task: _owned Builtin.NativeObject,
+///     options: Builtin.RawPointer?
+/// ) async -> Success
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_task_future_wait(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *,
+         SWIFT_ASYNC_CONTEXT AsyncContext *,
+         AsyncTask *,
+         TaskOptionRecord *options,
          TaskContinuationFunction *,
          AsyncContext *);
 
@@ -144,14 +148,17 @@ void swift_task_future_wait(OpaqueValue *,
 /// This can be called from any thread. Its Swift signature is
 ///
 /// \code
-/// func swift_task_future_wait_throwing(on task: _owned Builtin.NativeObject)
-///    async throws -> Success
+/// func swift_task_future_wait_throwing(
+///     on task: _owned Builtin.NativeObject,
+///     options: Builtin.RawPointer?
+/// ) async throws -> Success
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_task_future_wait_throwing(
   OpaqueValue *,
   SWIFT_ASYNC_CONTEXT AsyncContext *,
   AsyncTask *,
+  TaskOptionRecord *options,
   ThrowingTaskFutureWaitContinuationFunction *,
   AsyncContext *);
 
@@ -161,16 +168,18 @@ void swift_task_future_wait_throwing(
 ///
 /// \code
 /// func swift_taskGroup_wait_next_throwing(
-///     waitingTask: Builtin.NativeObject, // current task
-///     group: Builtin.RawPointer
+///     group: Builtin.RawPointer,
+///     options: Builtin.RawPointer?
 /// ) async -> T
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swiftasync)
 void swift_taskGroup_wait_next_throwing(
-    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
-    TaskGroup *group, ThrowingTaskFutureWaitContinuationFunction *resumeFn,
-    AsyncContext *callContext);
+    OpaqueValue *resultPointer,
+    SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    TaskGroup *group,
+    TaskOptionRecord *options,
+    ThrowingTaskFutureWaitContinuationFunction *resumeFn, AsyncContext *callContext);
 
 /// Initialize a `TaskGroup` in the passed `group` memory location.
 /// The caller is responsible for retaining and managing the group's lifecycle.
@@ -285,7 +294,8 @@ void swift_asyncLet_start(AsyncLet *alet,
 using AsyncLetWaitSignature =
     SWIFT_CC(swiftasync)
     void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
+        TaskOptionRecord *options,
+        SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
 
 /// Wait for a non-throwing async-let to complete.
 ///
@@ -294,27 +304,31 @@ using AsyncLetWaitSignature =
 /// \code
 /// func swift_asyncLet_wait(
 ///     _ asyncLet: _owned Builtin.RawPointer
+///     options: Builtin.RawPointer?
 /// ) async -> Success
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_asyncLet_wait(OpaqueValue *,
                          SWIFT_ASYNC_CONTEXT AsyncContext *,
-                         AsyncLet *, TaskContinuationFunction *,
-                         AsyncContext *);
+                         AsyncLet *,
+                         TaskOptionRecord *options,
+                         TaskContinuationFunction *, AsyncContext *);
 
 /// Wait for a potentially-throwing async-let to complete.
 ///
 /// This can be called from any thread. Its Swift signature is
 ///
 /// \code
 /// func swift_asyncLet_wait_throwing(
-///     _ asyncLet: _owned Builtin.RawPointer
+///     _ asyncLet: _owned Builtin.RawPointer,
+///     options: Builtin.RawPointer?
 /// ) async throws -> Success
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_asyncLet_wait_throwing(OpaqueValue *,
                                   SWIFT_ASYNC_CONTEXT AsyncContext *,
                                   AsyncLet *,
+                                  TaskOptionRecord *options,
                                   ThrowingTaskFutureWaitContinuationFunction *,
                                   AsyncContext *);
 

lib/SILGen/SILGenApply.cpp

@@ -5835,6 +5835,8 @@ ManagedValue SILGenFunction::emitCancelAsyncTask(
 
 void SILGenFunction::completeAsyncLetChildTask(
     PatternBindingDecl *patternBinding, unsigned index) {
+  auto &C = patternBinding->getASTContext();
+
   SILValue asyncLet;
   bool isThrowing;
   std::tie(asyncLet, isThrowing)= AsyncLetChildTasks[{patternBinding, index}];
@@ -5845,15 +5847,21 @@ void SILGenFunction::completeAsyncLetChildTask(
       ? SGM.getAsyncLetGetThrowing()
       : SGM.getAsyncLetGet();
 
+  /// No task options, so just pass nil
+  auto taskOptions = B.createManagedOptionalNone(
+      SILLocation(patternBinding),
+      SILType::getOptionalType(SILType::getRawPointerType(C)));
+
   // Get the result from the async-let future.
   Type replacementTypes[] = {childResultType};
   auto subs = SubstitutionMap::get(asyncLetGet->getGenericSignature(),
                                    replacementTypes,
                                    ArrayRef<ProtocolConformanceRef>{});
   RValue childResult = emitApplyOfLibraryIntrinsic(
       SILLocation(patternBinding), asyncLetGet, subs,
-      { ManagedValue::forTrivialObjectRValue(asyncLet) },
-      SGFContext());
+      { ManagedValue::forTrivialObjectRValue(asyncLet),
+        taskOptions
+      }, SGFContext());
 
   // Write the child result into the pattern variables.
   emitAssignToPatternVars(

lib/SILGen/SILGenDecl.cpp

@@ -1167,12 +1167,12 @@ void SILGenFunction::emitPatternBinding(PatternBindingDecl *PBD,
       // If we can statically detect some option needs to be passed, e.g.
       // an executor preference, we'd construct the appropriate option here and
       // pass it to the async let start.
-      auto options = B.createManagedOptionalNone(
+      auto taskOptions = B.createManagedOptionalNone(
           loc, SILType::getOptionalType(SILType::getRawPointerType(C)));
 
       alet = emitAsyncLetStart(
           loc,
-          options.forward(*this), // options is B.createManagedOptionalNone
+          taskOptions.forward(*this),
           init->getType(),
           emitRValue(init).getScalarValue()
         ).forward(*this);

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -110,19 +110,23 @@ OVERRIDE_TASK(task_create_group_future_common, AsyncTaskAndContext, , , ,
 OVERRIDE_TASK(task_future_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
               SWIFT_CC(swiftasync), swift::,
               (OpaqueValue *result,
-               SWIFT_ASYNC_CONTEXT AsyncContext *callerContext, AsyncTask *task,
+               SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+               AsyncTask *task,
+               TaskOptionRecord *options,
                TaskContinuationFunction *resumeFunction,
                AsyncContext *callContext),
-              (result, callerContext, task, resumeFunction, callContext))
+              (result, callerContext, task, options, resumeFunction, callContext))
 
 OVERRIDE_TASK(task_future_wait_throwing, void,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
               swift::,
               (OpaqueValue *result,
-               SWIFT_ASYNC_CONTEXT AsyncContext *callerContext, AsyncTask *task,
+               SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+               AsyncTask *task,
+               TaskOptionRecord *options,
                ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
                AsyncContext *callContext),
-              (result, callerContext, task, resumeFunction, callContext))
+              (result, callerContext, task, options, resumeFunction, callContext))
 
 OVERRIDE_TASK(continuation_resume, void, SWIFT_EXPORT_FROM(swift_Concurrency),
               SWIFT_CC(swift), swift::,
@@ -175,19 +179,21 @@ OVERRIDE_ASYNC_LET(asyncLet_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
                    SWIFT_CC(swiftasync), swift::,
                    (OpaqueValue *result,
                        SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
-                       AsyncLet *alet, TaskContinuationFunction *resumeFn,
-                       AsyncContext *callContext),
-                   (result, callerContext, alet, resumeFn, callContext))
+                       AsyncLet *alet,
+                       TaskOptionRecord *options,
+                       TaskContinuationFunction *resumeFn, AsyncContext *callContext),
+                   (result, callerContext, alet, options, resumeFn, callContext))
 
 OVERRIDE_ASYNC_LET(asyncLet_wait_throwing, void,
                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
                    swift::,
                    (OpaqueValue *result,
                        SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
                        AsyncLet *alet,
+                       TaskOptionRecord *options,
                        ThrowingTaskFutureWaitContinuationFunction *resume,
                        AsyncContext *callContext),
-                   (result, callerContext, alet, resume, callContext))
+                   (result, callerContext, alet, options, resume, callContext))
 
 OVERRIDE_ASYNC_LET(asyncLet_end, void,
                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
@@ -212,10 +218,11 @@ OVERRIDE_TASK_GROUP(taskGroup_wait_next_throwing, void,
                     (OpaqueValue *resultPointer,
                      SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
                      TaskGroup *_group,
+                     TaskOptionRecord *options,
                      ThrowingTaskFutureWaitContinuationFunction *resumeFn,
                      AsyncContext *callContext),
-                    (resultPointer, callerContext, _group, resumeFn,
-                    callContext))
+                    (resultPointer, callerContext, _group, options,
+                     resumeFn, callContext))
 
 OVERRIDE_TASK_GROUP(taskGroup_isEmpty, bool,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),

stdlib/public/Concurrency/AsyncLet.cpp

@@ -130,22 +130,24 @@ static void swift_asyncLet_startImpl(AsyncLet *alet,
 SWIFT_CC(swiftasync)
 static void swift_asyncLet_waitImpl(
     OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
-    AsyncLet *alet, TaskContinuationFunction *resumeFunction,
-    AsyncContext *callContext) {
+    AsyncLet *alet,
+    TaskOptionRecord *options,
+    TaskContinuationFunction *resumeFunction, AsyncContext *callContext) {
   auto task = alet->getTask();
-  swift_task_future_wait(result, callerContext, task, resumeFunction,
-                         callContext);
+  swift_task_future_wait(result, callerContext, task, options,
+                         resumeFunction, callContext);
 }
 
 SWIFT_CC(swiftasync)
 static void swift_asyncLet_wait_throwingImpl(
     OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
     AsyncLet *alet,
+    TaskOptionRecord *options,
     ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
     AsyncContext * callContext) {
   auto task = alet->getTask();
-  swift_task_future_wait_throwing(result, callerContext, task, resumeFunction,
-                                  callContext);
+  swift_task_future_wait_throwing(result, callerContext, task, options,
+                                  resumeFunction, callContext);
 }
 
 // =============================================================================

stdlib/public/Concurrency/AsyncLet.swift

@@ -27,12 +27,18 @@ public func _asyncLetStart<T>(
 /// Similar to _taskFutureGet but for AsyncLet
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_asyncLet_wait")
-public func _asyncLetGet<T>(asyncLet: Builtin.RawPointer) async -> T
+public func _asyncLetGet<T>(
+  asyncLet: Builtin.RawPointer,
+  options: Builtin.RawPointer?
+) async -> T
 
 ///// Similar to _taskFutureGetThrowing but for AsyncLet
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_asyncLet_wait_throwing")
-public func _asyncLetGetThrowing<T>(asyncLet: Builtin.RawPointer) async throws -> T
+public func _asyncLetGetThrowing<T>(
+  asyncLet: Builtin.RawPointer,
+  options: Builtin.RawPointer?
+) async throws -> T
 
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_asyncLet_end")

stdlib/public/Concurrency/Task.cpp

@@ -724,13 +724,16 @@ static void swift_task_future_waitImpl(
   OpaqueValue *result,
   SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
   AsyncTask *task,
+  TaskOptionRecord *options,
   TaskContinuationFunction *resumeFn,
   AsyncContext *callContext) {
   // Suspend the waiting task.
   auto waitingTask = swift_task_getCurrent();
   waitingTask->ResumeTask = task_future_wait_resume_adapter;
   waitingTask->ResumeContext = callContext;
 
+  // TODO: check `options` for ExecutorTaskOptionRecord and use it to resume if it was set.
+
   // Wait on the future.
   assert(task->isFuture());
 
@@ -755,8 +758,10 @@ static void swift_task_future_waitImpl(
 
 SWIFT_CC(swiftasync)
 void swift_task_future_wait_throwingImpl(
-    OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    OpaqueValue *result,
+    SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
     AsyncTask *task,
+    TaskOptionRecord *options,
     ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
     AsyncContext *callContext) {
   auto waitingTask = swift_task_getCurrent();

stdlib/public/Concurrency/TaskGroup.cpp

@@ -650,12 +650,15 @@ SWIFT_CC(swiftasync)
 static void swift_taskGroup_wait_next_throwingImpl(
     OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
     TaskGroup *_group,
+    TaskOptionRecord *options,
     ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
     AsyncContext *rawContext) {
   auto waitingTask = swift_task_getCurrent();
   waitingTask->ResumeTask = task_group_wait_resume_adapter;
   waitingTask->ResumeContext = rawContext;
 
+  // TODO: check `options` for ExecutorTaskOptionRecord and use it to resume if it was set.
+
   auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
   context->ResumeParent =
       reinterpret_cast<TaskContinuationFunction *>(resumeFunction);

stdlib/public/Concurrency/TaskGroup.swift

@@ -335,7 +335,7 @@ public struct TaskGroup<ChildTaskResult> {
   public mutating func next() async -> ChildTaskResult? {
     // try!-safe because this function only exists for Failure == Never,
     // and as such, it is impossible to spawn a throwing child task.
-    return try! await _taskGroupWaitNext(group: _group)
+    return try! await _taskGroupWaitNext(group: _group, options: nil)
   }
 
   /// Await all the remaining tasks on this group.
@@ -576,13 +576,13 @@ public struct ThrowingTaskGroup<ChildTaskResult, Failure: Error> {
   /// It is possible to directly rethrow such error out of a `withTaskGroup` body
   /// function's body, causing all remaining tasks to be implicitly cancelled.
   public mutating func next() async throws -> ChildTaskResult? {
-    return try await _taskGroupWaitNext(group: _group)
+    return try await _taskGroupWaitNext(group: _group, options: nil)
   }
 
   /// - SeeAlso: `next()`
   public mutating func nextResult() async throws -> Result<ChildTaskResult, Failure>? {
     do {
-      guard let success: ChildTaskResult = try await _taskGroupWaitNext(group: _group) else {
+      guard let success: ChildTaskResult = try await _taskGroupWaitNext(group: _group, options: nil) else {
         return nil
       }
 
@@ -777,7 +777,10 @@ func _taskGroupIsCancelled(group: Builtin.RawPointer) -> Bool
 
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_taskGroup_wait_next_throwing")
-func _taskGroupWaitNext<T>(group: Builtin.RawPointer) async throws -> T?
+func _taskGroupWaitNext<T>(
+  group: Builtin.RawPointer,
+  options: Builtin.RawPointer?
+) async throws -> T?
 
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_task_hasTaskGroupStatusRecord")