[Associated type inference] Limit Failure inference to rethrows next()

When inferring a type witness for `AsyncIteratorProtocol` or
`AsyncSequence`'s `Failure` associated type, don't infer from a
generic parameter named `Failure`. Instead, use `next()` as a cue: if
it `rethrows`, use `Failure` from one of the conformances; if it
`throws`, use `any Error`. This is a more conservative inference rule,
and addresses a failure to infer a `Failure` type witness for some
fairly-obvious cases.

Fixes rdar://122514816.

Author: DougGregor

lib/Sema/AssociatedTypeInference.cpp

@@ -1632,13 +1632,15 @@ next_witness:;
   return result;
 }
 
-/// Determine whether this is AsyncIteratorProtocol.Failure associated type.
+/// Determine whether this is AsyncIteratorProtocol.Failure or
+/// AsyncSequenceProtoco.Failure associated type.
 static bool isAsyncIteratorProtocolFailure(AssociatedTypeDecl *assocType) {
   auto proto = assocType->getProtocol();
-  if (!proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol))
+  if (!proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol) &&
+      !proto->isSpecificProtocol(KnownProtocolKind::AsyncSequence))
     return false;
 
-  return assocType->getName().str().equals("Failure");
+  return assocType->getName() == assocType->getASTContext().Id_Failure;
 }
 
 /// Determine whether this is AsyncIteratorProtocol.next() function.
@@ -2149,23 +2151,11 @@ llvm::Optional<AbstractTypeWitness>
 AssociatedTypeInference::computeFailureTypeWitness(
     AssociatedTypeDecl *assocType,
     ArrayRef<std::pair<ValueDecl *, ValueDecl *>> valueWitnesses) const {
-  // Inference only applies to AsyncIteratorProtocol.Failure.
+  // Inference only applies to AsyncIteratorProtocol.Failure and
+  // AsyncSequence.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) {
@@ -2179,8 +2169,6 @@ AssociatedTypeInference::computeFailureTypeWitness(
         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();
@@ -2206,7 +2194,8 @@ AssociatedTypeInference::computeFailureTypeWitness(
 llvm::Optional<AbstractTypeWitness>
 AssociatedTypeInference::computeDefaultTypeWitness(
     AssociatedTypeDecl *assocType) const {
-  // Ignore the default for AsyncIteratorProtocol.Failure
+  // Ignore the default for AsyncIteratorProtocol.Failure and
+  // AsyncSequence.Failure.
   if (isAsyncIteratorProtocolFailure(assocType))
     return llvm::None;
 
@@ -2304,13 +2293,14 @@ AssociatedTypeInference::computeAbstractTypeWitness(
   if (const auto &typeWitness = computeDefaultTypeWitness(assocType))
     return typeWitness;
 
+  // Ignore the default for AsyncIteratorProtocol.Failure and
+  // AsyncSequence.Failure. We use the next() function to do inference.
+  if (isAsyncIteratorProtocolFailure(assocType))
+    return llvm::None;
+
   // If there is a generic parameter of the named type, use that.
   if (auto genericSig = dc->getGenericSignatureOfContext()) {
-    bool wantAllGenericParams = isAsyncIteratorProtocolFailure(assocType);
-    auto genericParams = wantAllGenericParams 
-        ? genericSig.getGenericParams()
-        : genericSig.getInnermostGenericParams();
-    for (auto gp : genericParams) {
+    for (auto gp : genericSig.getInnermostGenericParams()) {
       // Packs cannot witness associated type requirements.
       if (gp->isParameterPack())
         continue;
@@ -2342,6 +2332,11 @@ void AssociatedTypeInference::collectAbstractTypeWitnesses(
   // through same-type requirements of protocols.
   if (auto genericSig = dc->getGenericSignatureOfContext()) {
     for (auto *const assocType : unresolvedAssocTypes) {
+      // Ignore the generic parameters for AsyncIteratorProtocol.Failure and
+      // AsyncSequence.Failure.
+      if (isAsyncIteratorProtocolFailure(assocType))
+        continue;
+
       for (auto *gp : genericSig.getInnermostGenericParams()) {
         // Packs cannot witness associated type requirements.
         if (gp->isParameterPack())

test/Concurrency/async_iterator_inference.swift

@@ -21,6 +21,18 @@ struct TS: AsyncSequence {
   func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
 }
 
+@available(SwiftStdlib 5.1, *)
+struct SpecificTS<F: Error>: AsyncSequence {
+  typealias Element = Int
+  typealias Failure = F
+  struct AsyncIterator: AsyncIteratorProtocol {
+    typealias Failure = F
+    mutating func next() async throws -> Int? { nil }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
+}
+
 @available(SwiftStdlib 5.1, *)
 struct GenericTS<Failure: Error>: AsyncSequence {
   typealias Element = Int
@@ -42,16 +54,33 @@ struct SequenceAdapter<Base: AsyncSequence>: AsyncSequence {
   func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
 }
 
+public struct NormalThrowingAsyncSequence<Element, Failure>: AsyncSequence {
+  private let iteratorMaker: () -> AsyncIterator
+
+  public struct AsyncIterator: AsyncIteratorProtocol {
+    let nextMaker: () async throws -> Element?
+    public mutating func next() async throws -> Element? {
+      try await nextMaker()
+    }
+  }
+
+  public func makeAsyncIterator() -> AsyncIterator {
+    iteratorMaker()
+  }
+}
+
+
 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) {
+func testAssocTypeInference(sf: S.Failure, tsf: TS.Failure, gtsf1: GenericTS<MyError>.Failure, adapter: SequenceAdapter<SpecificTS<MyError>>.Failure, ntas: NormalThrowingAsyncSequence<String, 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'}}
+  let _: Int = gtsf1 // expected-error{{cannot convert value of type 'GenericTS<MyError>.Failure' (aka 'any Error') to specified type 'Int'}}
+  let _: Int = adapter // expected-error{{cannot convert value of type 'SequenceAdapter<SpecificTS<MyError>>.Failure' (aka 'MyError') to specified type 'Int'}}
+  let _: Int = ntas // expected-error{{cannot convert value of type 'NormalThrowingAsyncSequence<String, MyError>.Failure' (aka 'any Error') to specified type 'Int'}}
 }
 
 
@@ -66,9 +95,9 @@ case boom
 
 
 @available(SwiftStdlib 5.1, *)
-func testMyError(s: GenericTS<MyError>, so: GenericTS<OtherError>) async throws(MyError) {
+func testMyError(s: SpecificTS<MyError>, so: SpecificTS<OtherError>) async throws(MyError) {
   for try await x in s { _ = x }
 
   for try await x in so { _ = x }
-  // expected-error@-1{{thrown expression type 'OtherError' cannot be converted to error type 'MyError'}}
+  // expected-error@-1{{thrown expression type 'SpecificTS<OtherError>.AsyncIterator.Failure' (aka 'OtherError') cannot be converted to error type 'MyError'}}
 }