Merge pull request #34916 from slavapestov/with-unsafe-continuation

Concurrency: Implement withUnsafe[Throwing]Continuation

Author: slavapestov

docs/ABI/Mangling.rst

@@ -498,6 +498,7 @@ Types
   type ::= 'BO'                              // Builtin.UnknownObject (no longer a distinct type, but still used for AnyObject)
   type ::= 'Bo'                              // Builtin.NativeObject
   type ::= 'Bp'                              // Builtin.RawPointer
+  type ::= 'Bc'                              // Builtin.RawUnsafeContinuation
   type ::= 'Bt'                              // Builtin.SILToken
   type ::= type 'Bv' NATURAL '_'             // Builtin.Vec<n>x<type>
   type ::= 'Bw'                              // Builtin.Word

include/swift/AST/ASTContext.h

@@ -751,6 +751,7 @@ class ASTContext final {
   const CanType TheNativeObjectType;      /// Builtin.NativeObject
   const CanType TheBridgeObjectType;      /// Builtin.BridgeObject
   const CanType TheRawPointerType;        /// Builtin.RawPointer
+  const CanType TheRawUnsafeContinuationType; /// Builtin.RawUnsafeContinuation
   const CanType TheUnsafeValueBufferType; /// Builtin.UnsafeValueBuffer
   const CanType TheSILTokenType;          /// Builtin.SILToken
   const CanType TheIntegerLiteralType;    /// Builtin.IntegerLiteralType

include/swift/AST/Builtins.def

@@ -504,6 +504,20 @@ BUILTIN_SIL_OPERATION(DifferentiableFunction, "differentiableFunction", Special)
 /// linearFunction
 BUILTIN_SIL_OPERATION(LinearFunction, "linearFunction", Special)
 
+/// withUnsafeContinuation<T> : (Builtin.RawUnsafeContinuation -> ()) async -> T
+///
+/// Unsafely capture the current continuation and pass it to the given
+/// function value. Returns a value of type T when the continuation is
+/// resumed.
+BUILTIN_SIL_OPERATION(WithUnsafeContinuation, "withUnsafeContinuation", Special)
+
+/// withUnsafeThrowingContinuation<T> : (Builtin.RawUnsafeContinuation -> ()) async throws -> T
+///
+/// Unsafely capture the current continuation and pass it to the given
+/// function value. Returns a value of type T or throws an error when
+/// the continuation is resumed.
+BUILTIN_SIL_OPERATION(WithUnsafeThrowingContinuation, "withUnsafeThrowingContinuation", Special)
+
 #undef BUILTIN_SIL_OPERATION
 
 // BUILTIN_RUNTIME_CALL - A call into a runtime function.

include/swift/AST/TypeDifferenceVisitor.h

@@ -113,6 +113,7 @@ class CanTypeDifferenceVisitor : public CanTypePairVisitor<Impl, bool> {
   }
   SINGLETON_TYPE(BuiltinIntegerLiteralType)
   SINGLETON_TYPE(BuiltinRawPointerType)
+  SINGLETON_TYPE(BuiltinRawUnsafeContinuationType)
   SINGLETON_TYPE(BuiltinNativeObjectType)
   SINGLETON_TYPE(BuiltinBridgeObjectType)
   SINGLETON_TYPE(BuiltinUnsafeValueBufferType)

include/swift/AST/TypeMatcher.h

@@ -105,6 +105,7 @@ class TypeMatcher {
     TRIVIAL_CASE(BuiltinIntegerLiteralType)
     TRIVIAL_CASE(BuiltinFloatType)
     TRIVIAL_CASE(BuiltinRawPointerType)
+    TRIVIAL_CASE(BuiltinRawUnsafeContinuationType)
     TRIVIAL_CASE(BuiltinNativeObjectType)
     TRIVIAL_CASE(BuiltinBridgeObjectType)
     TRIVIAL_CASE(BuiltinUnsafeValueBufferType)

include/swift/AST/TypeNodes.def

@@ -97,6 +97,7 @@ ABSTRACT_TYPE(Builtin, Type)
     TYPE_RANGE(AnyBuiltinInteger, BuiltinInteger, BuiltinIntegerLiteral)
   BUILTIN_TYPE(BuiltinFloat, BuiltinType)
   BUILTIN_TYPE(BuiltinRawPointer, BuiltinType)
+  BUILTIN_TYPE(BuiltinRawUnsafeContinuation, BuiltinType)
   BUILTIN_TYPE(BuiltinNativeObject, BuiltinType)
   BUILTIN_TYPE(BuiltinBridgeObject, BuiltinType)
   BUILTIN_TYPE(BuiltinUnsafeValueBuffer, BuiltinType)

include/swift/AST/Types.h

@@ -1391,6 +1391,20 @@ class BuiltinRawPointerType : public BuiltinType {
 };
 DEFINE_EMPTY_CAN_TYPE_WRAPPER(BuiltinRawPointerType, BuiltinType);
 
+/// BuiltinRawContinuationType - The builtin raw unsafe continuation type.
+/// This pointer is completely unmanaged, but is lowered with the same
+/// spare bits as an opaque object-pointer type.
+class BuiltinRawUnsafeContinuationType : public BuiltinType {
+  friend class ASTContext;
+  BuiltinRawUnsafeContinuationType(const ASTContext &C)
+    : BuiltinType(TypeKind::BuiltinRawUnsafeContinuation, C) {}
+public:
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::BuiltinRawUnsafeContinuation;
+  }
+};
+DEFINE_EMPTY_CAN_TYPE_WRAPPER(BuiltinRawUnsafeContinuationType, BuiltinType);
+
 /// BuiltinNativeObjectType - The builtin opaque object-pointer type.
 /// Useful for keeping an object alive when it is otherwise being
 /// manipulated via an unsafe pointer type.

include/swift/SIL/SILBuilder.h

@@ -1939,17 +1939,27 @@ class SILBuilder {
   //===--------------------------------------------------------------------===//
 
   GetAsyncContinuationInst *createGetAsyncContinuation(SILLocation Loc,
-                                                       SILType ContinuationTy) {
+                                                       CanType ResumeType,
+                                                       bool Throws) {
+    auto ContinuationType = SILType::getPrimitiveObjectType(
+        getASTContext().TheRawUnsafeContinuationType);
     return insert(new (getModule()) GetAsyncContinuationInst(getSILDebugLocation(Loc),
-                                                             ContinuationTy));
+                                                             ContinuationType,
+                                                             ResumeType,
+                                                             Throws));
   }
 
   GetAsyncContinuationAddrInst *createGetAsyncContinuationAddr(SILLocation Loc,
                                                                SILValue Operand,
-                                                               SILType ContinuationTy) {
+                                                               CanType ResumeType,
+                                                               bool Throws) {
+    auto ContinuationType = SILType::getPrimitiveObjectType(
+        getASTContext().TheRawUnsafeContinuationType);
     return insert(new (getModule()) GetAsyncContinuationAddrInst(getSILDebugLocation(Loc),
                                                                  Operand,
-                                                                 ContinuationTy));
+                                                                 ContinuationType,
+                                                                 ResumeType,
+                                                                 Throws));
   }
 
   HopToExecutorInst *createHopToExecutor(SILLocation Loc, SILValue Actor) {

include/swift/SIL/SILCloner.h

@@ -2961,7 +2961,8 @@ ::visitGetAsyncContinuationInst(GetAsyncContinuationInst *Inst) {
   recordClonedInstruction(Inst,
                           getBuilder().createGetAsyncContinuation(
                             getOpLocation(Inst->getLoc()),
-                            getOpType(Inst->getType())));
+                            getOpASTType(Inst->getFormalResumeType()),
+                            Inst->throws()));
 }
 
 template <typename ImplClass>
@@ -2972,7 +2973,8 @@ ::visitGetAsyncContinuationAddrInst(GetAsyncContinuationAddrInst *Inst) {
                           getBuilder().createGetAsyncContinuationAddr(
                             getOpLocation(Inst->getLoc()),
                             getOpValue(Inst->getOperand()),
-                            getOpType(Inst->getType())));
+                            getOpASTType(Inst->getFormalResumeType()),
+                            Inst->throws()));
 }
 
 template <typename ImplClass>

include/swift/SIL/SILInstruction.h

@@ -3244,15 +3244,22 @@ class GetAsyncContinuationInstBase
     : public SingleValueInstruction
 {
 protected:
-  using SingleValueInstruction::SingleValueInstruction;
-  
+  CanType ResumeType;
+  bool Throws;
+
+  GetAsyncContinuationInstBase(SILInstructionKind Kind, SILDebugLocation Loc,
+                               SILType ContinuationType, CanType ResumeType,
+                               bool Throws)
+    : SingleValueInstruction(Kind, Loc, ContinuationType),
+      ResumeType(ResumeType), Throws(Throws) {}
+
 public:
   /// Get the type of the value the async task receives on a resume.
-  CanType getFormalResumeType() const;
+  CanType getFormalResumeType() const { return ResumeType; }
   SILType getLoweredResumeType() const;
 
   /// True if the continuation can be used to resume the task by throwing an error.
-  bool throws() const;
+  bool throws() const { return Throws; }
 
   static bool classof(const SILNode *I) {
     return I->getKind() >= SILNodeKind::First_GetAsyncContinuationInstBase &&
@@ -3268,8 +3275,9 @@ class GetAsyncContinuationInst final
   friend SILBuilder;
 
   GetAsyncContinuationInst(SILDebugLocation Loc,
-                           SILType ContinuationTy)
-    : InstructionBase(Loc, ContinuationTy)
+                           SILType ContinuationType, CanType ResumeType,
+                           bool Throws)
+    : InstructionBase(Loc, ContinuationType, ResumeType, Throws)
   {}
 
 public:
@@ -3289,9 +3297,10 @@ class GetAsyncContinuationAddrInst final
 {
   friend SILBuilder;
   GetAsyncContinuationAddrInst(SILDebugLocation Loc,
-                               SILValue Operand,
-                               SILType ContinuationTy)
-    : UnaryInstructionBase(Loc, Operand, ContinuationTy)
+                               SILValue ResumeBuf,
+                               SILType ContinuationType, CanType ResumeType,
+                               bool Throws)
+    : UnaryInstructionBase(Loc, ResumeBuf, ContinuationType, ResumeType, Throws)
   {}
 };