Merge pull request #36621 from rjmccall/continuation-abi

Revise the ABI for interacting with unsafe continuations

Author: rjmccall

docs/ABI/Mangling.rst

@@ -517,6 +517,7 @@ Types
   type ::= 'BB'                              // Builtin.UnsafeValueBuffer
   type ::= 'Bc'                              // Builtin.RawUnsafeContinuation
   type ::= 'BD'                              // Builtin.DefaultActorStorage
+  type ::= 'Be'                              // Builtin.ExecutorRef
   type ::= 'Bf' NATURAL '_'                  // Builtin.Float<n>
   type ::= 'Bi' NATURAL '_'                  // Builtin.Int<n>
   type ::= 'BI'                              // Builtin.IntLiteral

include/swift/ABI/MetadataValues.h

@@ -2052,6 +2052,9 @@ enum class AsyncContextKind {
   /// A context which can yield to its caller.
   Yielding         = 1,
 
+  /// A continuation context.
+  Continuation     = 2,
+
   // Other kinds are reserved for interesting special
   // intermediate contexts.
 
@@ -2096,6 +2099,55 @@ class AsyncContextFlags : public FlagSet<uint32_t> {
                                 setShouldNotDeallocateInCallee)
 };
 
+/// Flags passed to swift_continuation_init.
+class AsyncContinuationFlags : public FlagSet<size_t> {
+public:
+  enum {
+    CanThrow            = 0,
+    HasExecutorOverride = 1,
+    IsPreawaited        = 2,
+  };
+
+  explicit AsyncContinuationFlags(size_t bits) : FlagSet(bits) {}
+  constexpr AsyncContinuationFlags() {}
+
+  /// Whether the continuation is permitted to throw.
+  FLAGSET_DEFINE_FLAG_ACCESSORS(CanThrow, canThrow, setCanThrow)
+
+  /// Whether the continuation should be resumed on a different
+  /// executor than the current one.  swift_continuation_init
+  /// will not initialize ResumeToExecutor if this is set.
+  FLAGSET_DEFINE_FLAG_ACCESSORS(HasExecutorOverride,
+                                hasExecutorOverride,
+                                setHasExecutorOverride)
+
+  /// Whether the continuation is "pre-awaited".  If so, it should
+  /// be set up in the already-awaited state, and so resumptions
+  /// will immediately schedule the continuation to begin
+  /// asynchronously.
+  FLAGSET_DEFINE_FLAG_ACCESSORS(IsPreawaited,
+                                isPreawaited,
+                                setIsPreawaited)
+};
+
+/// Status values for a continuation.  Note that the "not yet"s in
+/// the description below aren't quite right because the system
+/// does not actually promise to update the status before scheduling
+/// the task.  This is because the continuation context is immediately
+/// invalidated once the task starts running again, so the window in
+/// which we can usefully protect against (say) double-resumption may
+/// be very small.
+enum class ContinuationStatus : size_t {
+  /// The continuation has not yet been awaited or resumed.
+  Pending = 0,
+
+  /// The continuation has already been awaited, but not yet resumed.
+  Awaited = 1,
+
+  /// The continuation has already been resumed, but not yet awaited.
+  Resumed = 2
+};
+
 } // end namespace swift
 
 #endif // SWIFT_ABI_METADATAVALUES_H

include/swift/ABI/Task.h

@@ -536,6 +536,31 @@ class YieldingAsyncContext : public AsyncContext {
   }
 };
 
+/// An async context that can be resumed as a continuation.
+class ContinuationAsyncContext : public AsyncContext {
+public:
+  /// An atomic object used to ensure that a continuation is not
+  /// scheduled immediately during a resume if it hasn't yet been
+  /// awaited by the function which set it up.
+  std::atomic<ContinuationStatus> AwaitSynchronization;
+
+  /// The error result value of the continuation.
+  /// This should be null-initialized when setting up the continuation.
+  /// Throwing resumers must overwrite this with a non-null value.
+  SwiftError *ErrorResult;
+
+  /// A pointer to the normal result value of the continuation.
+  /// Normal resumers must initialize this before resuming.
+  OpaqueValue *NormalResult;
+
+  /// The executor that should be resumed to.
+  ExecutorRef ResumeToExecutor;
+
+  static bool classof(const AsyncContext *context) {
+    return context->Flags.getKind() == AsyncContextKind::Continuation;
+  }
+};
+
 /// An asynchronous context within a task that describes a general "Future".
 /// task.
 ///

include/swift/AST/ASTSynthesis.h

@@ -41,10 +41,12 @@ enum SingletonTypeSynthesizer {
   _any,
   _bridgeObject,
   _error,
+  _executor,
   _job,
   _nativeObject,
   _never,
   _rawPointer,
+  _rawUnsafeContinuation,
   _void,
   _word,
 };
@@ -54,10 +56,12 @@ inline Type synthesizeType(SynthesisContext &SC,
   case _any: return SC.Context.TheAnyType;
   case _bridgeObject: return SC.Context.TheBridgeObjectType;
   case _error: return SC.Context.getExceptionType();
+  case _executor: return SC.Context.TheExecutorType;
   case _job: return SC.Context.TheJobType;
   case _nativeObject: return SC.Context.TheNativeObjectType;
   case _never: return SC.Context.getNeverType();
   case _rawPointer: return SC.Context.TheRawPointerType;
+  case _rawUnsafeContinuation: return SC.Context.TheRawUnsafeContinuationType;
   case _void: return SC.Context.TheEmptyTupleType;
   case _word: return BuiltinIntegerType::get(BuiltinIntegerWidth::pointer(),
                                              SC.Context);

include/swift/AST/Builtins.def

@@ -723,6 +723,18 @@ BUILTIN_MISC_OPERATION(InitializeDefaultActor, "initializeDefaultActor", "", Spe
 /// Destroy the default-actor instance in a default actor object.
 BUILTIN_MISC_OPERATION(DestroyDefaultActor, "destroyDefaultActor", "", Special)
 
+/// Resume a non-throwing continuation normally with the given result.
+BUILTIN_MISC_OPERATION(ResumeNonThrowingContinuationReturning,
+                       "resumeNonThrowingContinuationReturning", "", Special)
+
+/// Resume a throwing continuation normally with the given result.
+BUILTIN_MISC_OPERATION(ResumeThrowingContinuationReturning,
+                       "resumeThrowingContinuationReturning", "", Special)
+
+/// Resume a throwing continuation abnormally with the given error.
+BUILTIN_MISC_OPERATION(ResumeThrowingContinuationThrowing,
+                       "resumeThrowingContinuationThrowing", "", Special)
+
 // BUILTIN_MISC_OPERATION_WITH_SILGEN - Miscellaneous operations that are
 // specially emitted during SIL generation.
 //
@@ -736,7 +748,7 @@ BUILTIN_MISC_OPERATION(DestroyDefaultActor, "destroyDefaultActor", "", Special)
   BUILTIN_MISC_OPERATION(Id, Name, Attrs, Overload)
 #endif
 
-// getCurrentExecutor: () async -> Builtin.Word
+// getCurrentExecutor: () async -> Builtin.Executor
 //
 // Retrieve the ExecutorRef on which the current asynchronous
 // function is executing.

include/swift/AST/TypeNodes.def

@@ -103,6 +103,7 @@ ABSTRACT_TYPE(Builtin, Type)
     BUILTIN_TYPE(BuiltinInteger, AnyBuiltinIntegerType)
     BUILTIN_TYPE(BuiltinIntegerLiteral, AnyBuiltinIntegerType)
     TYPE_RANGE(AnyBuiltinInteger, BuiltinInteger, BuiltinIntegerLiteral)
+  BUILTIN_TYPE(BuiltinExecutor, BuiltinType)
   BUILTIN_TYPE(BuiltinFloat, BuiltinType)
   BUILTIN_TYPE(BuiltinJob, BuiltinType)
   BUILTIN_TYPE(BuiltinRawPointer, BuiltinType)
@@ -181,6 +182,7 @@ LAST_TYPE(Dictionary) // Sugared types are last to make isa<SugarType>() fast.
 #ifdef SINGLETON_TYPE
 SINGLETON_TYPE(IntegerLiteral, BuiltinIntegerLiteral)
 SINGLETON_TYPE(Job, BuiltinJob)
+SINGLETON_TYPE(Executor, BuiltinExecutor)
 SINGLETON_TYPE(RawPointer, BuiltinRawPointer)
 SINGLETON_TYPE(RawUnsafeContinuation, BuiltinRawUnsafeContinuation)
 SINGLETON_TYPE(NativeObject, BuiltinNativeObject)

include/swift/AST/Types.h

@@ -1416,6 +1416,19 @@ class BuiltinRawUnsafeContinuationType : public BuiltinType {
 };
 DEFINE_EMPTY_CAN_TYPE_WRAPPER(BuiltinRawUnsafeContinuationType, BuiltinType);
 
+/// BuiltinExecutorType - The builtin executor-ref type.  In C, this
+/// is the ExecutorRef struct type.
+class BuiltinExecutorType : public BuiltinType {
+  friend class ASTContext;
+  BuiltinExecutorType(const ASTContext &C)
+    : BuiltinType(TypeKind::BuiltinExecutor, C) {}
+public:
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::BuiltinExecutor;
+  }
+};
+DEFINE_EMPTY_CAN_TYPE_WRAPPER(BuiltinExecutorType, BuiltinType);
+
 /// BuiltinJobType - The builtin job type.  In C, this is a
 /// non-null Job*.  This pointer is completely unmanaged (the unscheduled
 /// job is self-owning), but has more spare bits than Builtin.RawPointer.

include/swift/Runtime/Concurrency.h

@@ -523,23 +523,35 @@ void swift_defaultActor_destroy(DefaultActor *actor);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_defaultActor_enqueue(Job *job, DefaultActor *actor);
 
-/// Resume a task from its continuation, given a normal result value.
+/// Prepare a continuation in the current task.
+///
+/// The caller should initialize the Parent, ResumeParent,
+/// and NormalResult fields.  This function will initialize the other
+/// fields with appropriate defaaults; the caller may then overwrite
+/// them if desired.
+///
+/// This function is provided as a code-size and runtime-usage
+/// optimization; calling it is not required if code is willing to
+/// do all its work inline.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+AsyncTask *swift_continuation_init(ContinuationAsyncContext *context,
+                                   AsyncContinuationFlags flags);
+
+/// Resume a task from a non-throwing continuation, given a normal
+/// result which has already been stored into the continuation.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_continuation_resume(/* +1 */ OpaqueValue *result,
-                               void *continuation,
-                               const Metadata *resumeType);
+void swift_continuation_resume(AsyncTask *continuation);
 
-/// Resume a task from its throwing continuation, given a normal result value.
+/// Resume a task from a potentially-throwing continuation, given a
+/// normal result which has already been stored into the continuation.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_continuation_throwingResume(/* +1 */ OpaqueValue *result,
-                                       void *continuation,
-                                       const Metadata *resumeType);
+void swift_continuation_throwingResume(AsyncTask *continuation);
 
-/// Resume a task from its throwing continuation by throwing an error.
+/// Resume a task from a potentially-throwing continuation by throwing
+/// an error.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_continuation_throwingResumeWithError(/* +1 */ SwiftError *error,
-                                                void *continuation,
-                                                const Metadata *resumeType);
+void swift_continuation_throwingResumeWithError(AsyncTask *continuation,
+                                                /* +1 */ SwiftError *error);
 
 /// SPI helper to log a misuse of a `CheckedContinuation` to the appropriate places in the OS.
 extern "C" SWIFT_CC(swift)

include/swift/Runtime/RuntimeFunctions.def

@@ -1581,6 +1581,40 @@ FUNCTION(TaskSwitchFunc,
          ARGS(SwiftContextPtrTy, Int8PtrTy, SwiftExecutorPtrTy),
          ATTRS(NoUnwind))
 
+// AsyncTask *swift_continuation_init(AsyncContext *continuationContext,
+//                                    AsyncContinuationFlags);
+FUNCTION(ContinuationInit,
+         swift_continuation_init, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(SwiftTaskPtrTy),
+         ARGS(ContinuationAsyncContextPtrTy, SizeTy),
+         ATTRS(NoUnwind))
+
+// void swift_continuation_resume(AsyncTask *continuation);
+FUNCTION(ContinuationResume,
+         swift_continuation_resume, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(SwiftTaskPtrTy),
+         ATTRS(NoUnwind))
+
+// void swift_continuation_throwingResume(AsyncTask *continuation);
+FUNCTION(ContinuationThrowingResume,
+         swift_continuation_throwingResume, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(SwiftTaskPtrTy),
+         ATTRS(NoUnwind))
+
+// void swift_continuation_throwingResumeWithError(AsyncTask *continuation,
+//                                                 SwiftError *error);
+FUNCTION(ContinuationThrowingResumeWithError,
+         swift_continuation_throwingResumeWithError, SwiftCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(SwiftTaskPtrTy, ErrorPtrTy),
+         ATTRS(NoUnwind))
+
 // ExecutorRef swift_task_getCurrentExecutor();
 FUNCTION(TaskGetCurrentExecutor,
          swift_task_getCurrentExecutor, SwiftCC,

include/swift/Strings.h

@@ -122,6 +122,9 @@ constexpr static BuiltinNameStringLiteral BUILTIN_TYPE_NAME_UNSAFEVALUEBUFFER =
 /// The name of the Builtin type for Job
 constexpr static BuiltinNameStringLiteral BUILTIN_TYPE_NAME_JOB = {
     "Builtin.Job"};
+/// The name of the Builtin type for ExecutorRef
+constexpr static BuiltinNameStringLiteral BUILTIN_TYPE_NAME_EXECUTOR = {
+    "Builtin.ExecutorRef"};
 /// The name of the Builtin type for DefaultActorStorage
 constexpr static BuiltinNameStringLiteral BUILTIN_TYPE_NAME_DEFAULTACTORSTORAGE = {
     "Builtin.DefaultActorStorage"};

lib/AST/ASTDumper.cpp

@@ -3577,6 +3577,7 @@ namespace {
 
     TRIVIAL_TYPE_PRINTER(BuiltinIntegerLiteral, builtin_integer_literal)
     TRIVIAL_TYPE_PRINTER(BuiltinJob, builtin_job)
+    TRIVIAL_TYPE_PRINTER(BuiltinExecutor, builtin_executor_ref)
     TRIVIAL_TYPE_PRINTER(BuiltinDefaultActorStorage, builtin_default_actor_storage)
     TRIVIAL_TYPE_PRINTER(BuiltinRawPointer, builtin_raw_pointer)
     TRIVIAL_TYPE_PRINTER(BuiltinRawUnsafeContinuation, builtin_raw_unsafe_continuation)

lib/AST/ASTMangler.cpp

@@ -1047,6 +1047,8 @@ void ASTMangler::appendType(Type type, const ValueDecl *forDecl) {
       return appendOperator("BI");
     case TypeKind::BuiltinJob:
       return appendOperator("Bj");
+    case TypeKind::BuiltinExecutor:
+      return appendOperator("Be");
     case TypeKind::BuiltinDefaultActorStorage:
       return appendOperator("BD");
     case TypeKind::BuiltinRawPointer:

lib/AST/ASTPrinter.cpp

@@ -4407,6 +4407,7 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
   ASTPRINTER_PRINT_BUILTINTYPE(BuiltinRawPointerType)
   ASTPRINTER_PRINT_BUILTINTYPE(BuiltinRawUnsafeContinuationType)
   ASTPRINTER_PRINT_BUILTINTYPE(BuiltinJobType)
+  ASTPRINTER_PRINT_BUILTINTYPE(BuiltinExecutorType)
   ASTPRINTER_PRINT_BUILTINTYPE(BuiltinDefaultActorStorageType)
   ASTPRINTER_PRINT_BUILTINTYPE(BuiltinNativeObjectType)
   ASTPRINTER_PRINT_BUILTINTYPE(BuiltinBridgeObjectType)

lib/AST/Builtins.cpp

@@ -1351,10 +1351,9 @@ static ValueDecl *getGetCurrentAsyncTask(ASTContext &ctx, Identifier id) {
 }
 
 static ValueDecl *getGetCurrentExecutor(ASTContext &ctx, Identifier id) {
-  BuiltinFunctionBuilder builder(ctx);
-  builder.setResult(makeConcrete(BuiltinIntegerType::getWordType(ctx)));
-  builder.setAsync();
-  return builder.build(id);
+  return getBuiltinFunction(ctx, id, _async(_thin),
+                            _parameters(),
+                            _executor);
 }
 
 static ValueDecl *getCancelAsyncTask(ASTContext &ctx, Identifier id) {
@@ -1411,6 +1410,23 @@ static ValueDecl *getDefaultActorInitDestroy(ASTContext &ctx,
                             _void);
 }
 
+static ValueDecl *getResumeContinuationReturning(ASTContext &ctx,
+                                                 Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin,
+                            _generics(_unrestricted),
+                            _parameters(_rawUnsafeContinuation,
+                                        _owned(_typeparam(0))),
+                            _void);
+}
+
+static ValueDecl *getResumeContinuationThrowing(ASTContext &ctx,
+                                                Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin,
+                            _parameters(_rawUnsafeContinuation,
+                                        _owned(_error)),
+                            _void);
+}
+
 static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx,
                                                     Identifier id) {
   return getBuiltinFunction(
@@ -2621,6 +2637,13 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
   case BuiltinValueKind::DestroyDefaultActor:
     return getDefaultActorInitDestroy(Context, Id);
 
+  case BuiltinValueKind::ResumeNonThrowingContinuationReturning:
+  case BuiltinValueKind::ResumeThrowingContinuationReturning:
+    return getResumeContinuationReturning(Context, Id);
+
+  case BuiltinValueKind::ResumeThrowingContinuationThrowing:
+    return getResumeContinuationThrowing(Context, Id);
+
   case BuiltinValueKind::WithUnsafeContinuation:
     return getWithUnsafeContinuation(Context, Id, /*throws=*/false);
 
@@ -2692,6 +2715,9 @@ StringRef BuiltinType::getTypeName(SmallVectorImpl<char> &result,
   case BuiltinTypeKind::BuiltinJob:
     printer << MAYBE_GET_NAMESPACED_BUILTIN(BUILTIN_TYPE_NAME_JOB);
     break;
+  case BuiltinTypeKind::BuiltinExecutor:
+    printer << MAYBE_GET_NAMESPACED_BUILTIN(BUILTIN_TYPE_NAME_EXECUTOR);
+    break;
   case BuiltinTypeKind::BuiltinDefaultActorStorage:
     printer << MAYBE_GET_NAMESPACED_BUILTIN(BUILTIN_TYPE_NAME_DEFAULTACTORSTORAGE);
     break;

lib/AST/Type.cpp

@@ -212,6 +212,7 @@ bool CanType::isReferenceTypeImpl(CanType type, const GenericSignatureImpl *sig,
   case TypeKind::BuiltinRawPointer:
   case TypeKind::BuiltinRawUnsafeContinuation:
   case TypeKind::BuiltinJob:
+  case TypeKind::BuiltinExecutor:
   case TypeKind::BuiltinDefaultActorStorage:
   case TypeKind::BuiltinUnsafeValueBuffer:
   case TypeKind::BuiltinVector:
@@ -5097,6 +5098,7 @@ ReferenceCounting TypeBase::getReferenceCounting() {
   case TypeKind::BuiltinRawPointer:
   case TypeKind::BuiltinRawUnsafeContinuation:
   case TypeKind::BuiltinJob:
+  case TypeKind::BuiltinExecutor:
   case TypeKind::BuiltinDefaultActorStorage:
   case TypeKind::BuiltinUnsafeValueBuffer:
   case TypeKind::BuiltinVector: