Replace AsyncIteratorProtocol.nextElement() with isolated next(_:)

Use an optional isolated parameter to this new `next(_:)` overload to
keep it on the same actor as the caller, and pass `#isolation` when
desugaring the async for..in loop. This keeps async iteration loops on
the same actor, allowing non-Sendable values to be used with many
async sequences.

Author: DougGregor

include/swift/AST/ASTContext.h

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

include/swift/AST/KnownIdentifiers.def

@@ -123,7 +123,6 @@ IDENTIFIER(load)
 IDENTIFIER(main)
 IDENTIFIER_WITH_NAME(MainEntryPoint, "$main")
 IDENTIFIER(next)
-IDENTIFIER(nextElement)
 IDENTIFIER_(nsErrorDomain)
 IDENTIFIER(objectAtIndexedSubscript)
 IDENTIFIER(objectForKeyedSubscript)

lib/AST/ASTContext.cpp

@@ -287,8 +287,9 @@ struct ASTContext::Implementation {
   /// The declaration of 'AsyncIteratorProtocol.next()'.
   FuncDecl *AsyncIteratorNext = nullptr;
 
-  /// The declaration of 'AsyncIteratorProtocol.nextElement()'.
-  FuncDecl *AsyncIteratorNextElement = nullptr;
+  /// The declaration of 'AsyncIteratorProtocol.next(_:)' that takes
+  /// an actor isolation.
+  FuncDecl *AsyncIteratorNextIsolated = nullptr;
 
   /// The declaration of Swift.Optional<T>.Some.
   EnumElementDecl *OptionalSomeDecl = nullptr;
@@ -951,38 +952,48 @@ FuncDecl *ASTContext::getIteratorNext() const {
   return nullptr;
 }
 
-FuncDecl *ASTContext::getAsyncIteratorNext() const {
-  if (getImpl().AsyncIteratorNext) {
-    return getImpl().AsyncIteratorNext;
-  }
-
-  auto proto = getProtocol(KnownProtocolKind::AsyncIteratorProtocol);
+static std::pair<FuncDecl *, FuncDecl *>
+getAsyncIteratorNextRequirements(const ASTContext &ctx) {
+  auto proto = ctx.getProtocol(KnownProtocolKind::AsyncIteratorProtocol);
   if (!proto)
-    return nullptr;
+    return { nullptr, nullptr };
 
-  if (auto *func = lookupRequirement(proto, Id_next)) {
-    getImpl().AsyncIteratorNext = func;
-    return func;
+  FuncDecl *next = nullptr;
+  FuncDecl *nextThrowing = nullptr;
+  for (auto result : proto->lookupDirect(ctx.Id_next)) {
+    if (result->getDeclContext() != proto)
+      continue;
+
+    if (auto func = dyn_cast<FuncDecl>(result)) {
+      switch (func->getParameters()->size()) {
+      case 0: next = func; break;
+      case 1: nextThrowing = func; break;
+      default: break;
+      }
+    }
   }
 
-  return nullptr;
+  return { next, nextThrowing };
 }
 
-FuncDecl *ASTContext::getAsyncIteratorNextElement() const {
-  if (getImpl().AsyncIteratorNextElement) {
-    return getImpl().AsyncIteratorNextElement;
+FuncDecl *ASTContext::getAsyncIteratorNext() const {
+  if (getImpl().AsyncIteratorNext) {
+    return getImpl().AsyncIteratorNext;
   }
 
-  auto proto = getProtocol(KnownProtocolKind::AsyncIteratorProtocol);
-  if (!proto)
-    return nullptr;
+  auto next = getAsyncIteratorNextRequirements(*this).first;
+  getImpl().AsyncIteratorNext = next;
+  return next;
+}
 
-  if (auto *func = lookupRequirement(proto, Id_nextElement)) {
-    getImpl().AsyncIteratorNextElement = func;
-    return func;
+FuncDecl *ASTContext::getAsyncIteratorNextIsolated() const {
+  if (getImpl().AsyncIteratorNextIsolated) {
+    return getImpl().AsyncIteratorNextIsolated;
   }
 
-  return nullptr;
+  auto nextThrowing = getAsyncIteratorNextRequirements(*this).second;
+  getImpl().AsyncIteratorNextIsolated = nextThrowing;
+  return nextThrowing;
 }
 
 #define KNOWN_STDLIB_TYPE_DECL(NAME, DECL_CLASS, NUM_GENERIC_PARAMS) \

lib/Sema/CSGen.cpp

@@ -4603,7 +4603,15 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
         nextId, /*labels=*/ArrayRef<Identifier>());
     nextRef->setFunctionRefKind(FunctionRefKind::SingleApply);
 
-    Expr *nextCall = CallExpr::createImplicitEmpty(ctx, nextRef);
+    ArgumentList *nextArgs;
+    if (nextFn && nextFn->getParameters()->size() == 1) {
+      auto isolationArg =
+        new (ctx) CurrentContextIsolationExpr(stmt->getForLoc(), Type());
+      nextArgs = ArgumentList::forImplicitUnlabeled(ctx, { isolationArg });
+    } else {
+      nextArgs = ArgumentList::createImplicit(ctx, {});
+    }
+    Expr *nextCall = CallExpr::createImplicit(ctx, nextRef, nextArgs);
 
     // `next` is always async but witness might not be throwing
     if (isAsync) {

lib/Sema/CSSimplify.cpp

@@ -10705,7 +10705,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
       // Handle `next` reference.
       if (getContextualTypePurpose(baseExpr) == CTP_ForEachSequence &&
           (isRefTo(memberRef, ctx.Id_next, /*labels=*/{}) ||
-           isRefTo(memberRef, ctx.Id_nextElement, /*labels=*/{}))) {
+           isRefTo(memberRef, ctx.Id_next, /*labels=*/{StringRef()}))) {
         auto *iteratorProto = cast<ProtocolDecl>(
             getContextualType(baseExpr, /*forConstraint=*/false)
                 ->getAnyNominal());
@@ -10931,8 +10931,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
             if (auto *base = dyn_cast<DeclRefExpr>(UDE->getBase())) {
               if (auto var = dyn_cast_or_null<VarDecl>(base->getDecl())) {
                 if (var->getNameStr().contains("$generator") &&
-                    (UDE->getName().getBaseIdentifier() == Context.Id_next ||
-                     UDE->getName().getBaseIdentifier() == Context.Id_nextElement))
+                    (UDE->getName().getBaseIdentifier() == Context.Id_next))
                   return success();
               }
             }

lib/Sema/TypeCheckEffects.cpp

@@ -1335,7 +1335,7 @@ class ApplyClassifier {
             isRethrowLikeTypedThrows(fnRef.getFunction())) {
           // If we are in a rethrowing context and the function we're referring
           // to is a rethrow-like function using typed throws or we are
-          // calling the next() or nextElement() of an async iterator,
+          // calling the next() or next(_:) of an async iterator,
           // then look at all of the closure arguments.
           LLVM_FALLTHROUGH;
         } else {

lib/Sema/TypeCheckStmt.cpp

@@ -3237,17 +3237,17 @@ FuncDecl *TypeChecker::getForEachIteratorNextFunction(
   if (!isAsync)
     return ctx.getIteratorNext();
 
-  // If AsyncIteratorProtocol.nextElement() isn't available at all,
+  // If AsyncIteratorProtocol.next(_:) isn't available at all,
   // we're stuck using AsyncIteratorProtocol.next().
-  auto nextElement = ctx.getAsyncIteratorNextElement();
+  auto nextElement = ctx.getAsyncIteratorNextIsolated();
   if (!nextElement)
     return ctx.getAsyncIteratorNext();
 
-  // If availability checking is disabled, use nextElement().
+  // If availability checking is disabled, use next(_:).
   if (ctx.LangOpts.DisableAvailabilityChecking || loc.isInvalid())
     return nextElement;
 
-  // We can only call nextElement() if we are in an availability context
+  // We can only call next(_:) if we are in an availability context
   // that supports typed throws.
   auto availability = overApproximateAvailabilityAtLocation(loc, dc);
   if (availability.isContainedIn(ctx.getTypedThrowsAvailability()))

stdlib/public/Concurrency/AsyncCompactMapSequence.swift

@@ -140,9 +140,9 @@ extension AsyncCompactMapSequence: AsyncSequence {
     /// that transforms to a non-`nil` value.
     @available(SwiftStdlib 5.11, *)
     @inlinable
-    public mutating func nextElement() async throws(Failure) -> ElementOfResult? {
+    public mutating func next(_ actor: isolated (any Actor)?) async throws(Failure) -> ElementOfResult? {
       while true {
-        guard let element = try await baseIterator.nextElement() else {
+        guard let element = try await baseIterator.next(actor) else {
           return nil
         }
 

stdlib/public/Concurrency/AsyncDropFirstSequence.swift

@@ -118,24 +118,24 @@ extension AsyncDropFirstSequence: AsyncSequence {
     /// Produces the next element in the drop-first sequence.
     ///
     /// Until reaching the number of elements to drop, this iterator calls
-    /// `nextElement()` on its base iterator and discards the result. If the
+    /// `next(_:)` on its base iterator and discards the result. If the
     /// base iterator returns `nil`, indicating the end of the sequence, this
     /// iterator returns `nil`. After reaching the number of elements to drop,
-    /// this iterator passes along the result of calling `nextElement()` on the
+    /// this iterator passes along the result of calling `next(_:)` on the
     /// base iterator.
     @available(SwiftStdlib 5.11, *)
     @inlinable
-    public mutating func nextElement() async throws(Failure) -> Base.Element? {
+    public mutating func next(_ actor: isolated (any Actor)?) async throws(Failure) -> Base.Element? {
       var remainingToDrop = count
       while remainingToDrop > 0 {
-        guard try await baseIterator.nextElement() != nil else {
+        guard try await baseIterator.next(actor) != nil else {
           count = 0
           return nil
         }
         remainingToDrop -= 1
       }
       count = 0
-      return try await baseIterator.nextElement()
+      return try await baseIterator.next(actor)
     }
   }
 

stdlib/public/Concurrency/AsyncDropWhileSequence.swift

@@ -129,7 +129,7 @@ extension AsyncDropWhileSequence: AsyncSequence {
 
     /// Produces the next element in the drop-while sequence.
     ///
-    /// This iterator calls `nextElement()` on its base iterator and evaluates
+    /// This iterator calls `next(_:)` on its base iterator and evaluates
     /// the result with the `predicate` closure. As long as the predicate
     /// returns `true`, this method returns `nil`. After the predicate returns
     /// `false`, for a value received from the base iterator, this method
@@ -138,17 +138,17 @@ extension AsyncDropWhileSequence: AsyncSequence {
     /// again.
     @available(SwiftStdlib 5.11, *)
     @inlinable
-    public mutating func nextElement() async throws(Failure) -> Base.Element? {
+    public mutating func next(_ actor: isolated (any Actor)?) async throws(Failure) -> Base.Element? {
       while let predicate = self.predicate {
-        guard let element = try await baseIterator.nextElement() else {
+        guard let element = try await baseIterator.next(actor) else {
           return nil
         }
         if await predicate(element) == false {
           self.predicate = nil
           return element
         }
       }
-      return try await baseIterator.nextElement()
+      return try await baseIterator.next(actor)
     }
   }