Revise the continuation ABI.

The immediate desire is to minimize the set of ABI dependencies
on the layout of an ExecutorRef.  In addition to that, however,
I wanted to generally reduce the code size impact of an unsafe
continuation since it now requires accessing thread-local state,
and I wanted resumption to not have to create unnecessary type
metadata for the value type just to do the initialization.

Therefore, I've introduced a swift_continuation_init function
which handles the default initialization of a continuation
and returns a reference to the current task.  I've also moved
the initialization of the normal continuation result into the
caller (out of the runtime), and I've moved the resumption-side
cmpxchg into the runtime (and prior to the task being enqueued).

Author: rjmccall

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

@@ -46,6 +46,7 @@ enum SingletonTypeSynthesizer {
   _nativeObject,
   _never,
   _rawPointer,
+  _rawUnsafeContinuation,
   _void,
   _word,
 };
@@ -60,6 +61,7 @@ inline Type synthesizeType(SynthesisContext &SC,
   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.
 //

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,

lib/AST/Builtins.cpp

@@ -1410,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(
@@ -2620,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);
 

lib/IRGen/GenBuiltin.cpp

@@ -287,7 +287,23 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
     return;
   }
 
-  if (Builtin.ID == BuiltinValueKind::DestroyDefaultActor) {
+  if (Builtin.ID == BuiltinValueKind::ResumeThrowingContinuationReturning ||
+      Builtin.ID == BuiltinValueKind::ResumeNonThrowingContinuationReturning) {
+    auto continuation = args.claimNext();
+    auto valueTy = argTypes[1];
+    auto valuePtr = args.claimNext();
+    bool throwing =
+      (Builtin.ID == BuiltinValueKind::ResumeThrowingContinuationReturning);
+    IGF.emitResumeAsyncContinuationReturning(continuation, valuePtr, valueTy,
+                                             throwing);
+    return;
+  }
+
+  if (Builtin.ID == BuiltinValueKind::ResumeThrowingContinuationThrowing) {
+    auto continuation = args.claimNext();
+    auto error = args.claimNext();
+    IGF.emitResumeAsyncContinuationThrowing(continuation, error);
+    return;
   }
 
   // If this is an LLVM IR intrinsic, lower it to an intrinsic call.

lib/IRGen/GenCall.cpp

@@ -198,8 +198,8 @@ AsyncContextLayout::AsyncContextLayout(
 #endif
 }
 
-static Alignment getAsyncContextAlignment(IRGenModule &IGM) {
-  return IGM.getPointerAlignment();
+Alignment IRGenModule::getAsyncContextAlignment() const {
+  return Alignment(MaximumAlignment);
 }
 
 void IRGenFunction::setupAsync(unsigned asyncContextIndex) {
@@ -3815,7 +3815,7 @@ llvm::Value *irgen::emitTaskCreate(
 
 Address irgen::emitAllocAsyncContext(IRGenFunction &IGF,
                                      llvm::Value *sizeValue) {
-  auto alignment = getAsyncContextAlignment(IGF.IGM);
+  auto alignment = IGF.IGM.getAsyncContextAlignment();
   auto address = IGF.emitTaskAlloc(sizeValue, alignment);
   IGF.Builder.CreateLifetimeStart(address, Size(-1) /*dynamic size*/);
   return address;