[WIP] Typed throws for with(Unsafe|Checked)ThrowingContinuation

include/swift/AST/Builtins.def

@@ -524,7 +524,7 @@ BUILTIN_SIL_OPERATION(ApplyTranspose, "applyTranspose", Special)
 /// resumed.
 BUILTIN_SIL_OPERATION(WithUnsafeContinuation, "withUnsafeContinuation", Special)
 
-/// withUnsafeThrowingContinuation<T> : (Builtin.RawUnsafeContinuation -> ()) async throws -> sending T
+/// withUnsafeThrowingContinuation<T, E> : (Builtin.RawUnsafeContinuation -> ()) async throws(E) -> sending 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

lib/AST/Builtins.cpp

@@ -2137,27 +2137,31 @@ static ValueDecl *getPolymorphicBinaryOperation(ASTContext &ctx,
 static ValueDecl *getWithUnsafeContinuation(ASTContext &ctx,
                                             Identifier id,
                                             bool throws) {
-  BuiltinFunctionBuilder builder(ctx);
-
-  auto contTy = ctx.TheRawUnsafeContinuationType;
-  SmallVector<AnyFunctionType::Param, 1> params;
-  params.emplace_back(contTy);
-
-  auto voidTy = ctx.TheEmptyTupleType;
-  auto extInfo = FunctionType::ExtInfoBuilder().withNoEscape().build();
-  auto *fnTy = FunctionType::get(params, voidTy, extInfo);
-
-  builder.addParameter(makeConcrete(fnTy));
+  BuiltinFunctionBuilder builder(ctx, throws ? 2 : 1);
 
   auto resultTy = makeGenericParam();
   builder.addConformanceRequirement(resultTy, KnownProtocolKind::Escapable);
   builder.setResult(resultTy);
-
   builder.setAsync();
-  if (throws)
+
+  SmallVector<AnyFunctionType::Param, 1> params;
+  if (throws) {
+    auto errorTy = makeGenericParam(1);
+    builder.addConformanceRequirement(errorTy, KnownProtocolKind::Error);
     builder.setThrows();
+    builder.setThrownError(errorTy);
+    builder.addParameter(makeMetatype(errorTy));
+  }
   builder.setSendingResult();
 
+  // Add the closure parameter.
+  auto voidTy = ctx.TheEmptyTupleType;
+  auto contTy = ctx.TheRawUnsafeContinuationType;
+  params.emplace_back(contTy);
+  auto extInfo = FunctionType::ExtInfoBuilder().withNoEscape().build();
+  auto *fnTy = FunctionType::get(params, voidTy, extInfo);
+  builder.addParameter(makeConcrete(fnTy));
+
   return builder.build(id);
 }
 

lib/SILGen/SILGenBuiltin.cpp

@@ -1848,10 +1848,11 @@ static ManagedValue emitBuiltinWithUnsafeContinuation(
                                                            throws);
 
   // Get the callee value.
-  auto substFnType = args[0].getType().castTo<SILFunctionType>();
+  unsigned calleeIndex = throws ? 1 : 0;
+  auto substFnType = args[calleeIndex].getType().castTo<SILFunctionType>();
   SILValue fnValue = (substFnType->isCalleeConsumed()
-                      ? args[0].forward(SGF)
-                      : args[0].getValue());
+                      ? args[calleeIndex].forward(SGF)
+                      : args[calleeIndex].getValue());
 
   // Call the provided function value.
   SGF.B.createApply(loc, fnValue, {}, {continuation});
@@ -1871,7 +1872,13 @@ static ManagedValue emitBuiltinWithUnsafeContinuation(
 
     Scope errorScope(SGF, loc);
 
-    auto errorTy = SGF.getASTContext().getErrorExistentialType();
+    CanType errorTy;
+    if (subs.getReplacementTypes().size() > 1) {
+      errorTy = subs.getReplacementTypes()[1]->getCanonicalType();
+    } else {
+      errorTy = SGF.getASTContext().getErrorExistentialType();
+    }
+
     auto errorVal = SGF.B.createTermResult(
         SILType::getPrimitiveObjectType(errorTy), OwnershipKind::Owned);
 

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -357,14 +357,31 @@ public func _unsafeInheritExecutor_withCheckedContinuation<T>(
 @inlinable
 @available(SwiftStdlib 5.1, *)
 #if !$Embedded
-@backDeployed(before: SwiftStdlib 6.0)
+@backDeployed(before: SwiftStdlib 6.2)
 #endif
-public func withCheckedThrowingContinuation<T>(
+public func withCheckedThrowingContinuation<T, E>(
+  isolation: isolated (any Actor)? = #isolation,
+  function: String = #function,
+  _ body: (CheckedContinuation<T, E>) -> Void
+) async throws(E) -> sending T {
+  return try await Builtin.withUnsafeThrowingContinuation(E.self) {
+    let unsafeContinuation = unsafe UnsafeContinuation<T, E>($0)
+    return body(unsafe CheckedContinuation(continuation: unsafeContinuation,
+                                           function: function))
+  }
+}
+
+// Superseded by the typed-throws version of this function. This function
+// is retained for ABI purposes.
+@available(SwiftStdlib 5.1, *)
+@usableFromInline
+@_silgen_name("$ss31withCheckedThrowingContinuation9isolation8function_xScA_pSgYi_SSyScCyxs5Error_pGXEtYaKlF")
+internal func __abi_withCheckedThrowingContinuation<T>(
   isolation: isolated (any Actor)? = #isolation,
   function: String = #function,
   _ body: (CheckedContinuation<T, Error>) -> Void
 ) async throws -> sending T {
-  return try await Builtin.withUnsafeThrowingContinuation {
+  return try await Builtin.withUnsafeThrowingContinuation(Error.self) {
     let unsafeContinuation = unsafe UnsafeContinuation<T, Error>($0)
     return body(unsafe CheckedContinuation(continuation: unsafeContinuation,
                                            function: function))

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -721,12 +721,12 @@ public func withUnsafeContinuation<T>(
 @available(SwiftStdlib 5.1, *)
 @_alwaysEmitIntoClient
 @unsafe
-public func withUnsafeThrowingContinuation<T>(
+public func withUnsafeThrowingContinuation<T, E>(
   isolation: isolated (any Actor)? = #isolation,
-  _ fn: (UnsafeContinuation<T, Error>) -> Void
-) async throws -> sending T {
-  return try await Builtin.withUnsafeThrowingContinuation {
-    unsafe fn(UnsafeContinuation<T, Error>($0))
+  _ fn: (UnsafeContinuation<T, E>) -> Void
+) async throws(E) -> sending T {
+  return try await Builtin.withUnsafeThrowingContinuation(E.self) {
+    unsafe fn(UnsafeContinuation<T, E>($0))
   }
 }
 
@@ -750,11 +750,11 @@ public func _unsafeInheritExecutor_withUnsafeContinuation<T>(
 @available(SwiftStdlib 5.1, *)
 @_alwaysEmitIntoClient
 @_unsafeInheritExecutor
-public func _unsafeInheritExecutor_withUnsafeThrowingContinuation<T>(
-  _ fn: (UnsafeContinuation<T, Error>) -> Void
-) async throws -> sending T {
-  return try await Builtin.withUnsafeThrowingContinuation {
-    unsafe fn(UnsafeContinuation<T, Error>($0))
+public func _unsafeInheritExecutor_withUnsafeThrowingContinuation<T, E>(
+  _ fn: (UnsafeContinuation<T, E>) -> Void
+) async throws(E) -> sending T {
+  return try await Builtin.withUnsafeThrowingContinuation(E.self) {
+    unsafe fn(UnsafeContinuation<T, E>($0))
   }
 }
 

test/Concurrency/with_continuation_typed_throws.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-frontend -target %target-swift-5.1-abi-triple %s -typecheck /dev/null -verify
+
+// REQUIRES: concurrency
+
+enum MyError: Error {
+case exploded
+}
+
+func testTypedThrowsContinuations() async throws(MyError) { 
+  let _: Int = try await withUnsafeThrowingContinuation { (continuation: UnsafeContinuation<Int, MyError>) in
+    continuation.resume(throwing: .exploded)
+  }
+
+  let _: Int = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Int, MyError>) in
+    continuation.resume(throwing: .exploded)
+  }
+}