Adopt typed throws in AsyncIteratorProtocol and AsyncSequence

Introduce a new associated type `Failure` into the two protocols involved
in async sequences, which represents the type thrown when the sequence
fails. Introduce a defaulted `_nextElement()` operations that throws
`Failure` or produces the next element of the sequence. Provide a
default implementation of `_nextElement()` in terms of `next()` that
force-cases the thrown error to the `Failure` type.

Introduce special associated type inference logic for the `Failure`
type of an `AsyncIteratorProtocol` conformance when there is no
specific _nextElement()` witness. This inference logic looks at the
witness for `next()`:
* If `next()` throws nothing, `Failure` is inferred to `Never`.
* If `next()` throws, `Failure` is inferred to `any Error`.
* If `next()` rethrows, `Failure` is inferred to `T.Failure`, where
`T` is the first type parameter with a conformance to either
`AsyncSequence` or `AsyncIteratorProtocol`.

The default implementation and the inference rule, together, allow
existing async sequences to continue working as before, and set us up
for changing the contract of the `async for` loop to use
`_nextIterator()` rather than `next()`.

Give `AsyncSequence` and `AsyncIteratorProtocol` primary associated
types for the element and failure types, which will allow them to be
used more generally with existential and opaque types.

Author: DougGregor

lib/Sema/AssociatedTypeInference.cpp

@@ -821,6 +821,14 @@ class swift::AssociatedTypeInference {
   llvm::Optional<AbstractTypeWitness>
   computeDefaultTypeWitness(AssociatedTypeDecl *assocType) const;
 
+  /// Compute type witnesses for the Failure type from the
+  /// AsyncSequence or AsyncIteratorProtocol
+  llvm::Optional<AbstractTypeWitness>
+  computeFailureTypeWitness(
+      AssociatedTypeDecl *assocType,
+      ArrayRef<std::pair<ValueDecl *, ValueDecl *>> valueWitnesses
+  ) const;
+
   /// Compute the "derived" type witness for an associated type that is
   /// known to the compiler.
   std::pair<Type, TypeDecl *>
@@ -1370,6 +1378,25 @@ next_witness:;
   return result;
 }
 
+/// Determine whether this is AsyncIteratorProtocol.Failure associated type.
+static bool isAsyncIteratorProtocolFailure(AssociatedTypeDecl *assocType) {
+  auto proto = assocType->getProtocol();
+  if (!proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol))
+    return false;
+
+  return assocType->getName().str().equals("Failure");
+}
+
+/// Determine whether this is AsyncIteratorProtocol.next() function.
+static bool isAsyncIteratorProtocolNext(ValueDecl *req) {
+  auto proto = dyn_cast<ProtocolDecl>(req->getDeclContext());
+  if (!proto ||
+      !proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol))
+    return false;
+
+  return req->getName().getBaseName() == req->getASTContext().Id_next;
+}
+
 InferredAssociatedTypes
 AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
   const llvm::SetVector<AssociatedTypeDecl *> &assocTypes) {
@@ -1415,7 +1442,8 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
                                           TinyPtrVector<AssociatedTypeDecl *>());
       if (llvm::find_if(referenced, [&](AssociatedTypeDecl *const assocType) {
                           return assocTypes.count(assocType);
-                        }) == referenced.end())
+                        }) == referenced.end() &&
+          !isAsyncIteratorProtocolNext(req))
         continue;
     }
 
@@ -1845,9 +1873,71 @@ Type AssociatedTypeInference::computeFixedTypeWitness(
   return resultType;
 }
 
+llvm::Optional<AbstractTypeWitness>
+AssociatedTypeInference::computeFailureTypeWitness(
+    AssociatedTypeDecl *assocType,
+    ArrayRef<std::pair<ValueDecl *, ValueDecl *>> valueWitnesses) const {
+  // Inference only applies to AsyncIteratorProtocol.Failure.
+  if (!isAsyncIteratorProtocolFailure(assocType))
+    return llvm::None;
+
+  // If there is a generic parameter named Failure, don't try to use next()
+  // to infer Failure.
+  if (auto genericSig = dc->getGenericSignatureOfContext()) {
+    for (auto gp : genericSig.getGenericParams()) {
+      // Packs cannot witness associated type requirements.
+      if (gp->isParameterPack())
+        continue;
+
+      if (gp->getName() == assocType->getName())
+        return llvm::None;
+    }
+  }
+
+  // Look for AsyncIteratorProtocol.next() and infer the Failure type from
+  // it.
+  for (const auto &witness : valueWitnesses) {
+    if (isAsyncIteratorProtocolNext(witness.first)) {
+      if (auto witnessFunc = dyn_cast<AbstractFunctionDecl>(witness.second)) {
+        // If it doesn't throw, Failure == Never.
+        if (!witnessFunc->hasThrows())
+          return AbstractTypeWitness(assocType, ctx.getNeverType());
+
+        // If it isn't 'rethrows', Failure == any Error.
+        if (!witnessFunc->getAttrs().hasAttribute<RethrowsAttr>())
+          return AbstractTypeWitness(assocType, ctx.getErrorExistentialType());
+
+        // Otherwise, we need to derive the Failure type from a type parameter
+        // that conforms to AsyncIteratorProtocol or AsyncSequence.
+        for (auto req : witnessFunc->getGenericSignature().getRequirements()) {
+          if (req.getKind() == RequirementKind::Conformance) {
+            auto proto = req.getProtocolDecl();
+            if (proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol) ||
+                proto->isSpecificProtocol(KnownProtocolKind::AsyncSequence)) {
+              auto failureAssocType = proto->getAssociatedType(ctx.getIdentifier("Failure"));
+              auto failureType = DependentMemberType::get(req.getFirstType(), failureAssocType);
+              return AbstractTypeWitness(assocType, dc->mapTypeIntoContext(failureType));
+            }
+          }
+        }
+
+        return AbstractTypeWitness(assocType, ctx.getErrorExistentialType());
+      }
+
+      break;
+    }
+  }
+
+  return llvm::None;
+}
+
 llvm::Optional<AbstractTypeWitness>
 AssociatedTypeInference::computeDefaultTypeWitness(
     AssociatedTypeDecl *assocType) const {
+  // Ignore the default for AsyncIteratorProtocol.Failure
+  if (isAsyncIteratorProtocolFailure(assocType))
+    return llvm::None;
+
   // Go find a default definition.
   auto *const defaultedAssocType = findDefaultedAssociatedType(
       dc, dc->getSelfNominalTypeDecl(), assocType);
@@ -1942,7 +2032,11 @@ AssociatedTypeInference::computeAbstractTypeWitness(
 
   // If there is a generic parameter of the named type, use that.
   if (auto genericSig = dc->getGenericSignatureOfContext()) {
-    for (auto gp : genericSig.getInnermostGenericParams()) {
+    bool wantAllGenericParams = isAsyncIteratorProtocolFailure(assocType);
+    auto genericParams = wantAllGenericParams 
+        ? genericSig.getGenericParams()
+        : genericSig.getInnermostGenericParams();
+    for (auto gp : genericParams) {
       // Packs cannot witness associated type requirements.
       if (gp->isParameterPack())
         continue;
@@ -2552,7 +2646,20 @@ void AssociatedTypeInference::findSolutionsRec(
     // Filter out the associated types that remain unresolved.
     SmallVector<AssociatedTypeDecl *, 4> stillUnresolved;
     for (auto *const assocType : unresolvedAssocTypes) {
-      const auto typeWitness = typeWitnesses.begin(assocType);
+      auto typeWitness = typeWitnesses.begin(assocType);
+
+      // If we do not have a witness for AsyncIteratorProtocol.Failure,
+      // look for the witness to AsyncIteratorProtocol.next(). If it throws,
+      // use 'any Error'. Otherwise, use 'Never'.
+      if (typeWitness == typeWitnesses.end()) {
+        if (auto failureTypeWitness =
+                computeFailureTypeWitness(assocType, valueWitnesses)) {
+          typeWitnesses.insert(assocType,
+                               {failureTypeWitness->getType(), reqDepth});
+          typeWitness = typeWitnesses.begin(assocType);
+        }
+      }
+
       if (typeWitness == typeWitnesses.end()) {
         stillUnresolved.push_back(assocType);
       } else {

lib/Sema/TypeCheckEffects.cpp

@@ -171,6 +171,17 @@ static bool classifyWitness(ModuleDecl *module,
 
     case PolymorphicEffectKind::Always:
       // Witness always has the effect.
+
+      // If the witness's thrown type is explicitly specified as a type
+      // parameter, then check whether the substituted type is `Never`.
+      if (kind == EffectKind::Throws) {
+        if (Type thrownError = witnessDecl->getThrownInterfaceType()) {
+          if (thrownError->hasTypeParameter())
+            thrownError = thrownError.subst(declRef.getSubstitutions());
+          if (thrownError->isNever())
+            return false;
+        }
+      }
       return true;
 
     case PolymorphicEffectKind::Invalid:

stdlib/public/Concurrency/AsyncIteratorProtocol.swift

@@ -87,12 +87,38 @@ import Swift
 /// a reference type.
 @available(SwiftStdlib 5.1, *)
 @rethrows
-public protocol AsyncIteratorProtocol {
+public protocol AsyncIteratorProtocol<Element, Failure> {
   associatedtype Element
+
+  /// The type of failure produced by iteration.
+  associatedtype Failure: Error = any Error
+
   /// Asynchronously advances to the next element and returns it, or ends the
   /// sequence if there is no next element.
-  /// 
+  ///
   /// - Returns: The next element, if it exists, or `nil` to signal the end of
   ///   the sequence.
   mutating func next() async throws -> Element?
+
+  /// Asynchronously advances to the next element and returns it, or ends the
+  /// sequence if there is no next element.
+  ///
+  /// - Returns: The next element, if it exists, or `nil` to signal the end of
+  ///   the sequence.
+  @available(SwiftStdlib 5.11, *)
+  mutating func _nextElement() async throws(Failure) -> Element?
+}
+
+@available(SwiftStdlib 5.1, *)
+extension AsyncIteratorProtocol {
+  /// Default implementation of `_nextElement()` in terms of `next()`, which is
+  /// required to maintain backward compatibility with existing async iterators.
+  @available(SwiftStdlib 5.11, *)
+  public mutating func _nextElement() async throws(Failure) -> Element? {
+    do {
+      return try await next()
+    } catch {
+      throw error as! Failure
+    }
+  }
 }

stdlib/public/Concurrency/AsyncSequence.swift

@@ -73,12 +73,17 @@ import Swift
 ///
 @available(SwiftStdlib 5.1, *)
 @rethrows
-public protocol AsyncSequence {
+public protocol AsyncSequence<Element, Failure> {
   /// The type of asynchronous iterator that produces elements of this
   /// asynchronous sequence.
   associatedtype AsyncIterator: AsyncIteratorProtocol where AsyncIterator.Element == Element
   /// The type of element produced by this asynchronous sequence.
   associatedtype Element
+
+  /// The type of errors produced when iteration over the sequence fails.
+  associatedtype Failure: Error = AsyncIterator.Failure
+      where AsyncIterator.Failure == Failure
+
   /// Creates the asynchronous iterator that produces elements of this
   /// asynchronous sequence.
   ///

stdlib/public/Concurrency/CMakeLists.txt

@@ -169,7 +169,7 @@ add_swift_target_library(swift_Concurrency ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} I
   SWIFT_COMPILE_FLAGS
     ${SWIFT_STANDARD_LIBRARY_SWIFT_FLAGS}
     -parse-stdlib
-    -Xfrontend -enable-experimental-concurrency
+    -enable-experimental-feature TypedThrows
     -diagnostic-style swift
     ${SWIFT_RUNTIME_CONCURRENCY_SWIFT_FLAGS}
   ${swift_concurrency_options}
@@ -256,6 +256,7 @@ if(SWIFT_SHOULD_BUILD_EMBEDDED_STDLIB AND SWIFT_SHOULD_BUILD_EMBEDDED_CONCURRENC
 
       SWIFT_COMPILE_FLAGS
         ${extra_swift_compile_flags} -enable-experimental-feature Embedded
+        -enable-experimental-feature TypedThrows
         -parse-stdlib -DSWIFT_CONCURRENCY_EMBEDDED
         ${SWIFT_RUNTIME_CONCURRENCY_SWIFT_FLAGS}
       C_COMPILE_FLAGS

test/Concurrency/async_iterator_inference.swift

@@ -0,0 +1,61 @@
+// RUN: %target-swift-frontend -strict-concurrency=complete -emit-sil -o /dev/null %s -verify
+// REQUIRES: concurrency
+
+@available(SwiftStdlib 5.1, *)
+struct S: AsyncSequence {
+  typealias Element = Int
+  struct AsyncIterator: AsyncIteratorProtocol {
+    mutating func next() async -> Int? { nil }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
+}
+
+@available(SwiftStdlib 5.1, *)
+struct TS: AsyncSequence {
+  typealias Element = Int
+  struct AsyncIterator: AsyncIteratorProtocol {
+    mutating func next() async throws -> Int? { nil }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
+}
+
+@available(SwiftStdlib 5.1, *)
+struct GenericTS<Failure: Error>: AsyncSequence {
+  typealias Element = Int
+  struct AsyncIterator: AsyncIteratorProtocol {
+    mutating func next() async throws -> Int? { nil }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
+}
+
+@available(SwiftStdlib 5.1, *)
+struct SequenceAdapter<Base: AsyncSequence>: AsyncSequence {
+  typealias Element = Base.Element
+
+  struct AsyncIterator: AsyncIteratorProtocol {
+    mutating func next() async rethrows -> Base.Element? { nil }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
+}
+
+enum MyError: Error {
+case fail
+}
+
+@available(SwiftStdlib 5.1, *)
+func testAssocTypeInference(sf: S.Failure, tsf: TS.Failure, gtsf1: GenericTS<MyError>.Failure, adapter: SequenceAdapter<GenericTS<MyError>>.Failure) {
+  let _: Int = sf // expected-error{{cannot convert value of type 'S.Failure' (aka 'Never') to specified type 'Int'}}
+  let _: Int = tsf // expected-error{{cannot convert value of type 'TS.Failure' (aka 'any Error') to specified type 'Int'}}
+  let _: Int = gtsf1 // expected-error{{cannot convert value of type 'GenericTS<MyError>.Failure' (aka 'MyError') to specified type 'Int'}}
+  let _: Int = adapter // expected-error{{cannot convert value of type 'SequenceAdapter<GenericTS<MyError>>.Failure' (aka 'MyError') to specified type 'Int'}}
+}
+
+
+@available(SwiftStdlib 5.1, *)
+func test(s: S) async {
+  for await x in s { _ = x }
+}

test/api-digester/stability-concurrency-abi.test

@@ -60,6 +60,8 @@ Func _asyncLet_get_throwing(_:_:) has mangled name changing from '_Concurrency._
 Func _asyncLet_get_throwing(_:_:) has return type change from Builtin.RawPointer to ()
 Protocol Actor has added inherited protocol AnyActor
 Protocol Actor has generic signature change from <Self : AnyObject, Self : Swift.Sendable> to <Self : _Concurrency.AnyActor>
+Protocol AsyncIteratorProtocol has generic signature change from to <Self.Failure : Swift.Error>
+Protocol AsyncSequence has generic signature change from <Self.AsyncIterator : _Concurrency.AsyncIteratorProtocol, Self.Element == Self.AsyncIterator.Element> to <Self.AsyncIterator : _Concurrency.AsyncIteratorProtocol, Self.Element == Self.AsyncIterator.Element, Self.Failure == Self.AsyncIterator.Failure>
 Struct CheckedContinuation has removed conformance to UnsafeSendable
 
 // SerialExecutor gained `enqueue(_: __owned Job)`, protocol requirements got default implementations