Choose between AsyncIteratorProtocol's next() and nextElement() based on availability

This allows us to not break backward deployment

Author: DougGregor

include/swift/AST/ASTContext.h

@@ -650,6 +650,9 @@ class ASTContext final {
   /// Get AsyncIteratorProtocol.next().
   FuncDecl *getAsyncIteratorNext() const;
 
+  /// Get AsyncIteratorProtocol.nextElement().
+  FuncDecl *getAsyncIteratorNextElement() const;
+
   /// Check whether the standard library provides all the correct
   /// intrinsic support for Optional<T>.
   ///

lib/AST/ASTContext.cpp

@@ -287,6 +287,9 @@ struct ASTContext::Implementation {
   /// The declaration of 'AsyncIteratorProtocol.next()'.
   FuncDecl *AsyncIteratorNext = nullptr;
 
+  /// The declaration of 'AsyncIteratorProtocol.nextElement()'.
+  FuncDecl *AsyncIteratorNextElement = nullptr;
+
   /// The declaration of Swift.Optional<T>.Some.
   EnumElementDecl *OptionalSomeDecl = nullptr;
 
@@ -957,13 +960,25 @@ FuncDecl *ASTContext::getAsyncIteratorNext() const {
   if (!proto)
     return nullptr;
 
-  if (auto *func = lookupRequirement(proto, Id_nextElement)) {
+  if (auto *func = lookupRequirement(proto, Id_next)) {
     getImpl().AsyncIteratorNext = func;
     return func;
   }
 
-  if (auto *func = lookupRequirement(proto, Id_next)) {
-    getImpl().AsyncIteratorNext = func;
+  return nullptr;
+}
+
+FuncDecl *ASTContext::getAsyncIteratorNextElement() const {
+  if (getImpl().AsyncIteratorNextElement) {
+    return getImpl().AsyncIteratorNextElement;
+  }
+
+  auto proto = getProtocol(KnownProtocolKind::AsyncIteratorProtocol);
+  if (!proto)
+    return nullptr;
+
+  if (auto *func = lookupRequirement(proto, Id_nextElement)) {
+    getImpl().AsyncIteratorNextElement = func;
     return func;
   }
 

lib/AST/Availability.cpp

@@ -819,6 +819,7 @@ ASTContext::getSwift5PlusAvailability(llvm::VersionTuple swiftVersion) {
     case 7: return getSwift57Availability();
     case 8: return getSwift58Availability();
     case 9: return getSwift59Availability();
+    case 11: return getSwift511Availability();
     default: break;
     }
   }

lib/Sema/CSGen.cpp

@@ -4592,8 +4592,8 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
   // Now, result type of `.makeIterator()` is used to form a call to
   // `.next()`. `next()` is called on each iteration of the loop.
   {
-    FuncDecl *nextFn = isAsync ? ctx.getAsyncIteratorNext()
-                               : ctx.getIteratorNext();
+    FuncDecl *nextFn = 
+        TypeChecker::getForEachIteratorNextFunction(dc, stmt->getForLoc(), isAsync);
     Identifier nextId = nextFn ? nextFn->getName().getBaseIdentifier()
                                : ctx.Id_next;
     auto *nextRef = UnresolvedDotExpr::createImplicit(

lib/Sema/CSSimplify.cpp

@@ -10712,8 +10712,8 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
         bool isAsync = iteratorProto->getKnownProtocolKind() ==
                        KnownProtocolKind::AsyncIteratorProtocol;
 
-        auto *next =
-            isAsync ? ctx.getAsyncIteratorNext() : ctx.getIteratorNext();
+        auto loc = locator->getAnchor().getStartLoc();
+        auto *next = TypeChecker::getForEachIteratorNextFunction(DC, loc, isAsync);
 
         return simplifyValueWitnessConstraint(
             ConstraintKind::ValueWitness, baseTy, next, memberTy, useDC,

lib/Sema/TypeCheckStmt.cpp

@@ -3227,3 +3227,32 @@ void swift::bindSwitchCasePatternVars(DeclContext *dc, CaseStmt *caseStmt) {
     recordVar(nullptr, bodyVar);
   }
 }
+
+FuncDecl *TypeChecker::getForEachIteratorNextFunction(
+    DeclContext *dc, SourceLoc loc, bool isAsync
+) {
+  ASTContext &ctx = dc->getASTContext();
+
+  // A synchronous for..in loop uses IteratorProtocol.next().
+  if (!isAsync)
+    return ctx.getIteratorNext();
+
+  // If AsyncIteratorProtocol.nextElement() isn't available at all,
+  // we're stuck using AsyncIteratorProtocol.next().
+  auto nextElement = ctx.getAsyncIteratorNextElement();
+  if (!nextElement)
+    return ctx.getAsyncIteratorNext();
+
+  // If availability checking is disabled, use nextElement().
+  if (ctx.LangOpts.DisableAvailabilityChecking || loc.isInvalid())
+    return nextElement;
+
+  // We can only call nextElement() if we are in an availability context
+  // that supports typed throws.
+  auto availability = overApproximateAvailabilityAtLocation(loc, dc);
+  if (availability.isContainedIn(ctx.getTypedThrowsAvailability()))
+    return nextElement;
+
+  // Fall back to AsyncIteratorProtocol.next().
+  return ctx.getAsyncIteratorNext();
+}

lib/Sema/TypeChecker.h

@@ -1169,6 +1169,11 @@ Type catchErrorType(DeclContext *dc, DoCatchStmt *stmt);
 Type errorUnion(Type type1, Type type2,
                 llvm::function_ref<Type(Type)> simplifyType);
 
+/// Retrieve the "next" function that should be used for iteration in a
+/// for..in loop.
+FuncDecl *getForEachIteratorNextFunction(
+    DeclContext *dc, SourceLoc loc, bool isAsync);
+
 /// If an expression references 'self.init' or 'super.init' in an
 /// initializer context, returns the implicit 'self' decl of the constructor.
 /// Otherwise, return nil.

test/Concurrency/async_for_each_back_deploy.swift

@@ -0,0 +1,9 @@
+// RUN: %target-swift-frontend -strict-concurrency=complete -dump-ast %s | %FileCheck %s
+// REQUIRES: concurrency, OS=macosx
+
+@available(SwiftStdlib 5.1, *)
+func f<S: AsyncSequence>(s: S) async throws {
+  // CHECK-NOT: nextElement()
+  // CHECK: next()
+  for try await x in s { }
+}

test/Concurrency/async_iterator_inference.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -strict-concurrency=complete -emit-sil -o /dev/null %s -verify
+// RUN: %target-swift-frontend -strict-concurrency=complete -emit-sil -o /dev/null %s -verify -disable-availability-checking
 // REQUIRES: concurrency
 
 @available(SwiftStdlib 5.1, *)

test/IRGen/async/weak_availability.swift

@@ -9,5 +9,5 @@ public func f<S: AsyncSequence>(_ s: S) async throws -> Any.Type {
   return Fn.self
 }
 
-// MAYBE-AVAILABLE: @"$sScI11nextElement0B0QzSgyYa7FailureQzYKFTjTu" = extern_weak global
+// MAYBE-AVAILABLE: @"$sScI{{.*next.*}}YaKFTjTu" = extern_weak global
 // MAYBE-AVAILABLE: declare{{.*}} extern_weak{{.*}} @swift_getFunctionTypeMetadataGlobalActor