Merge pull request #77663 from rjmccall/self-consuming-create-task

Add support for creating a task with a self-consuming task function

Author: rjmccall

include/swift/ABI/MetadataValues.h

@@ -2654,6 +2654,13 @@ class TaskCreateFlags : public FlagSet<size_t> {
     Task_EnqueueJob                               = 12,
     Task_AddPendingGroupTaskUnconditionally       = 13,
     Task_IsDiscardingTask                         = 14,
+
+    /// The task function is consumed by calling it (@callee_owned).
+    /// The context pointer should be treated as opaque and non-copyable;
+    /// in particular, it should not be retained or released.
+    ///
+    /// Supported starting in Swift 6.1.
+    Task_IsTaskFunctionConsumed                       = 15,
   };
 
   explicit constexpr TaskCreateFlags(size_t bits) : FlagSet(bits) {}
@@ -2683,6 +2690,9 @@ class TaskCreateFlags : public FlagSet<size_t> {
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsDiscardingTask,
                                 isDiscardingTask,
                                 setIsDiscardingTask)
+  FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsTaskFunctionConsumed,
+                                isTaskFunctionConsumed,
+                                setIsTaskFunctionConsumed)
 };
 
 /// Flags for schedulable jobs.

include/swift/AST/FeatureAvailability.def

@@ -75,6 +75,7 @@ FEATURE(NoncopyableGenerics,                            (6, 0))
 FEATURE(InitRawStructMetadata,                          (6, 0))
 
 FEATURE(LayoutStringValueWitnesses,                     (6, 1))
+FEATURE(CreateTaskWithConsumedFunction,                 (6, 1))
 
 FEATURE(TaskExecutor,                                   FUTURE)
 FEATURE(Differentiation,                                FUTURE)

stdlib/public/Concurrency/Task.cpp

@@ -1044,16 +1044,44 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
                        " with parent %p at base pri %zu",
                        task, task->getTaskId(), parent, basePriority);
 
+  // Configure the initial context.
+
+  // Initialize the parent context pointer to null.
+  initialContext->Parent = nullptr;
+
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wcast-function-type-mismatch"
-  // Initialize the task-local allocator.
-  initialContext->ResumeParent =
-      runInlineOption ? &completeInlineTask
-                      : reinterpret_cast<TaskContinuationFunction *>(
-                            asyncLet         ? &completeTask
-                            : closureContext ? &completeTaskWithClosure
-                                             : &completeTaskAndRelease);
+  // Initialize the resumption funclet pointer (async return address) to
+  // the final funclet for completing the task.
+
+  // Inline tasks are unmanaged, non-throwing, and use a non-escaping
+  // task function.  The final funclet doesn't expect to get passed an error,
+  // and it doesn't clean up either the function or the task directly.
+  if (runInlineOption) {
+    initialContext->ResumeParent = &completeInlineTask;
+
+  // `async let` tasks are unmanaged and use a non-escaping task function.
+  // The final funclet shouldn't release the task or the task function.
+  } else if (asyncLet) {
+    initialContext->ResumeParent =
+      reinterpret_cast<TaskContinuationFunction*>(&completeTask);
+
+  // If we have a non-null closure context and the task function is not
+  // consumed by calling it, use a final funclet that releases both the
+  // task and the closure context.
+  } else if (closureContext && !taskCreateFlags.isTaskFunctionConsumed()) {
+    initialContext->ResumeParent =
+      reinterpret_cast<TaskContinuationFunction*>(&completeTaskWithClosure);
+
+  // Otherwise, just release the task.
+  } else {
+    initialContext->ResumeParent =
+      reinterpret_cast<TaskContinuationFunction*>(&completeTaskAndRelease);
+  }
 #pragma clang diagnostic pop
+
+  // Initialize the task-local allocator and our other private runtime
+  // state for the task.
   if ((asyncLet || (runInlineOption && runInlineOption->getAllocation())) &&
       initialSlabSize > 0) {
     assert(parent || (runInlineOption && runInlineOption->getAllocation()));
@@ -1095,14 +1123,6 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
     task->_private().Local.initializeLinkParent(task, parent);
   }
 
-  // Configure the initial context.
-  //
-  // FIXME: if we store a null pointer here using the standard ABI for
-  // signed null pointers, then we'll have to authenticate context pointers
-  // as if they might be null, even though the only time they ever might
-  // be is the final hop.  Store a signed null instead.
-  initialContext->Parent = nullptr;
-
   // FIXME: add discarding flag
   // FIXME: add task executor
   concurrency::trace::task_create(