Add nextElement() and explicit Failure types for a number of asynchronous iterators

Author: DougGregor

stdlib/public/Concurrency/AsyncStream.swift

@@ -377,6 +377,24 @@ extension AsyncStream: AsyncSequence {
     public mutating func next() async -> Element? {
       await context.produce()
     }
+
+    /// The next value from the asynchronous stream.
+    ///
+    /// When `nextElement()` returns `nil`, this signifies the end of the
+    /// `AsyncStream`.
+    ///
+    /// It is a programmer error to invoke `nextElement()` from a concurrent
+    /// context that contends with another such call, which results in a call to
+    /// `fatalError()`.
+    ///
+    /// If you cancel the task this iterator is running in while `nextElement()`
+    /// is awaiting a value, the `AsyncStream` terminates. In this case,
+    /// `nextElement()` might return `nil` immediately, or return `nil` on
+    /// subsequent calls.
+    @available(SwiftStdlib 5.11, *)
+    public mutating func nextElement() async -> Element? {
+      await context.produce()
+    }
   }
 
   /// Creates the asynchronous iterator that produces elements of this
@@ -530,6 +548,12 @@ extension AsyncStream {
     public mutating func next() async -> Element? {
       fatalError("Unavailable in task-to-thread concurrency model")
     }
+    
+    @available(SwiftStdlib 5.11, *)
+    @available(*, unavailable, message: "Unavailable in task-to-thread concurrency model")
+    public mutating func nextElement() async -> Element? {
+      fatalError("Unavailable in task-to-thread concurrency model")
+    }
   }
   @available(SwiftStdlib 5.1, *)
   @available(*, unavailable, message: "Unavailable in task-to-thread concurrency model")

stdlib/public/Concurrency/AsyncThrowingCompactMapSequence.swift

@@ -93,6 +93,11 @@ extension AsyncThrowingCompactMapSequence: AsyncSequence {
   /// The compact map sequence produces whatever type of element its
   /// transforming closure produces.
   public typealias Element = ElementOfResult
+  /// The type of element produced by this asynchronous sequence.
+  ///
+  /// The compact map sequence produces errors from either the base
+  /// sequence or the transforming closure.
+  public typealias Failure = any Error
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 
@@ -145,6 +150,35 @@ extension AsyncThrowingCompactMapSequence: AsyncSequence {
       }
       return nil
     }
+
+    /// Produces the next element in the compact map sequence.
+    ///
+    /// This iterator calls `nextElement()` on its base iterator; if this call
+    /// returns `nil`, `nextElement()` returns `nil`. Otherwise, `nextElement()`
+    /// calls the transforming closure on the received element, returning it if
+    /// the transform returns a non-`nil` value. If the transform returns `nil`,
+    /// this method continues to wait for further elements until it gets one
+    /// that transforms to a non-`nil` value. If calling the closure throws an
+    /// error, the sequence ends and `nextElement()` rethrows the error.
+    @available(SwiftStdlib 5.11, *)
+    @inlinable
+    public mutating func nextElement() async throws(Failure) -> ElementOfResult? {
+      while !finished {
+        guard let element = try await baseIterator.nextElement() else {
+          finished = true
+          return nil
+        }
+        do {
+          if let transformed = try await transform(element) {
+            return transformed
+          }
+        } catch {
+          finished = true
+          throw error
+        }
+      }
+      return nil
+    }
   }
 
   @inlinable

stdlib/public/Concurrency/AsyncThrowingDropWhileSequence.swift

@@ -91,6 +91,11 @@ extension AsyncThrowingDropWhileSequence: AsyncSequence {
   /// The drop-while sequence produces whatever type of element its base
   /// sequence produces.
   public typealias Element = Base.Element
+  /// The type of element produced by this asynchronous sequence.
+  ///
+  /// The drop-while sequence produces errors from either the base
+  /// sequence or the filtering closure.
+  public typealias Failure = any Error
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 
@@ -148,6 +153,39 @@ extension AsyncThrowingDropWhileSequence: AsyncSequence {
       }
       return try await baseIterator.next()
     }
+
+    /// Produces the next element in the drop-while sequence.
+    ///
+    /// This iterator calls `nextElement()` 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
+    /// returns that value. After that, the iterator returns values received
+    /// from its base iterator as-is, and never executes the predicate closure
+    /// again.  If calling the closure throws an error, the sequence ends and
+    /// `nextElement()` rethrows the error.
+    @available(SwiftStdlib 5.11, *)
+    @inlinable
+    public mutating func nextElement() async throws(Failure) -> Base.Element? {
+      while !finished && !doneDropping {
+        guard let element = try await baseIterator.nextElement() else {
+          return nil
+        }
+        do {
+          if try await predicate(element) == false {
+            doneDropping = true
+            return element
+          }
+        } catch {
+          finished = true
+          throw error
+        }
+      }
+      guard !finished else { 
+        return nil
+      }
+      return try await baseIterator.nextElement()
+    }
   }
 
   @inlinable

stdlib/public/Concurrency/AsyncThrowingFilterSequence.swift

@@ -81,6 +81,11 @@ extension AsyncThrowingFilterSequence: AsyncSequence {
   /// The filter sequence produces whatever type of element its base
   /// sequence produces.
   public typealias Element = Base.Element
+  /// The type of element produced by this asynchronous sequence.
+  ///
+  /// The filter sequence produces errors from either the base
+  /// sequence or the filtering closure.
+  public typealias Failure = any Error
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 
@@ -130,6 +135,35 @@ extension AsyncThrowingFilterSequence: AsyncSequence {
 
       return nil
     }
+
+    /// Produces the next element in the filter sequence.
+    ///
+    /// This iterator calls `nextElement()` on its base iterator; if this call
+    /// returns `nil`, `nextElement()` returns nil. Otherwise, `nextElement()`
+    /// evaluates the result with the `predicate` closure. If the closure
+    /// returns `true`, `nextElement()` returns the received element; otherwise
+    /// it awaits the next element from the base iterator. If calling the
+    /// closure throws an error, the sequence ends and `nextElement()` rethrows
+    /// the error.
+    @available(SwiftStdlib 5.11, *)
+    @inlinable
+    public mutating func nextElement() async throws(Failure) -> Base.Element? {
+      while !finished {
+        guard let element = try await baseIterator.nextElement() else {
+          return nil
+        }
+        do {
+          if try await isIncluded(element) {
+            return element
+          }
+        } catch {
+          finished = true
+          throw error
+        }
+      }
+
+      return nil
+    }
   }
 
   @inlinable

stdlib/public/Concurrency/AsyncThrowingFlatMapSequence.swift

@@ -89,6 +89,11 @@ extension AsyncThrowingFlatMapSequence: AsyncSequence {
   /// The flat map sequence produces the type of element in the asynchronous
   /// sequence produced by the `transform` closure.
   public typealias Element = SegmentOfResult.Element
+  /// The type of error produced by this asynchronous sequence.
+  ///
+  /// The flat map sequence produces errors from either the base
+  /// sequence or the `transform` closure.
+  public typealias Failure = any Error
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 

stdlib/public/Concurrency/AsyncThrowingMapSequence.swift

@@ -92,6 +92,11 @@ extension AsyncThrowingMapSequence: AsyncSequence {
   /// The map sequence produces whatever type of element its the transforming
   /// closure produces.
   public typealias Element = Transformed
+  /// The type of error produced by this asynchronous sequence.
+  ///
+  /// The map sequence produces errors from either the base
+  /// sequence or the `transform` closure.
+  public typealias Failure = any Error
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 
@@ -122,7 +127,7 @@ extension AsyncThrowingMapSequence: AsyncSequence {
     /// calling the transforming closure on the received element. If calling
     /// the closure throws an error, the sequence ends and `next()` rethrows
     /// the error.
-   @inlinable
+    @inlinable
     public mutating func next() async throws -> Transformed? {
       guard !finished, let element = try await baseIterator.next() else {
         return nil
@@ -134,6 +139,27 @@ extension AsyncThrowingMapSequence: AsyncSequence {
         throw error   
       }
     }
+
+    /// Produces the next element in the map sequence.
+    ///
+    /// This iterator calls `nextElement()` on its base iterator; if this call
+    /// returns `nil`, `nextElement()` returns nil. Otherwise, `nextElement()`
+    /// returns the result of calling the transforming closure on the received
+    /// element. If calling the closure throws an error, the sequence ends and
+    /// `nextElement()` rethrows the error.
+    @available(SwiftStdlib 5.11, *)
+    @inlinable
+    public mutating func nextElement() async throws(Failure) -> Transformed? {
+      guard !finished, let element = try await baseIterator.nextElement() else {
+        return nil
+      }
+      do {
+        return try await transform(element)
+      } catch {
+        finished = true
+        throw error   
+      }
+    }
   }
 
   @inlinable

stdlib/public/Concurrency/AsyncThrowingPrefixWhileSequence.swift

@@ -87,6 +87,11 @@ extension AsyncThrowingPrefixWhileSequence: AsyncSequence {
   /// The prefix-while sequence produces whatever type of element its base
   /// iterator produces.
   public typealias Element = Base.Element
+  /// The type of error produced by this asynchronous sequence.
+  ///
+  /// The prefix-while sequence produces errors from either the base
+  /// sequence or the filtering closure.
+  public typealias Failure = any Error
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 
@@ -133,6 +138,31 @@ extension AsyncThrowingPrefixWhileSequence: AsyncSequence {
       }
       return nil
     }
+
+    /// Produces the next element in the prefix-while sequence.
+    ///
+    /// If the predicate hasn't failed yet, this method gets the next element
+    /// from the base sequence and calls the predicate with it. If this call
+    /// succeeds, this method passes along the element. Otherwise, it returns
+    /// `nil`, ending the sequence. If calling the predicate closure throws an
+    /// error, the sequence ends and `nextElement()` rethrows the error.
+    @available(SwiftStdlib 5.11, *)
+    @inlinable
+    public mutating func nextElement() async throws(Failure) -> Base.Element? {
+      if !predicateHasFailed, let nextElement = try await baseIterator.nextElement() {
+        do { 
+          if try await predicate(nextElement) {
+            return nextElement
+          } else {
+            predicateHasFailed = true
+          }
+        } catch {
+          predicateHasFailed = true
+          throw error
+        }
+      }
+      return nil
+    }
   }
 
   @inlinable

stdlib/public/Concurrency/TaskGroup.swift

@@ -1141,6 +1141,36 @@ extension ThrowingTaskGroup: AsyncSequence {
       }
     }
 
+    /// Advances to and returns the result of the next child task.
+    ///
+    /// The elements returned from this method
+    /// appear in the order that the tasks *completed*,
+    /// not in the order that those tasks were added to the task group.
+    /// After this method returns `nil`,
+    /// this iterator is guaranteed to never produce more values.
+    ///
+    /// For more information about the iteration order and semantics,
+    /// see `ThrowingTaskGroup.next()`
+    ///
+    /// - Throws: The error thrown by the next child task that completes.
+    ///
+    /// - Returns: The value returned by the next child task that completes,
+    ///   or `nil` if there are no remaining child tasks,
+    @available(SwiftStdlib 5.11, *)
+    public mutating func nextElement() async throws(Failure) -> Element? {
+      guard !finished else { return nil }
+      do {
+        guard let element = try await group.next() else {
+          finished = true
+          return nil
+        }
+        return element
+      } catch {
+        finished = true
+        throw error as! Failure
+      }
+    }
+    
     public mutating func cancel() {
       finished = true
       group.cancelAll()

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

@@ -65,7 +65,6 @@ Protocol AsyncSequence has generic signature change from <Self.AsyncIterator : _
 Struct CheckedContinuation has removed conformance to UnsafeSendable
 Protocol AsyncIteratorProtocol is now without @rethrows
 Protocol AsyncSequence is now without @rethrows
-Func AsyncIteratorProtocol.nextElement() has been added as a protocol requirement
 
 // SerialExecutor gained `enqueue(_: __owned Job)`, protocol requirements got default implementations
 Func SerialExecutor.enqueue(_:) has been added as a protocol requirement