Adopt nextElement() in AsyncFlatMapSequence

`AsyncFlatMapSequence` is somewhat troublesome for typed throws,
because it can produce errors from two different sources: the `Base`
async sequence and the `SegmentOfResult` async sequence. However, we
need to pick one `Failure` type for the `AsyncFlatMapSequence`, and
there is no surface-language equivalent to the `errorUnion` operation
described in SE-0413.

So, we do something a little bit sneaky. We effectively require that
`SegmentOfResult.Failure` either be equivalent to `Base.Failure` or be
`Never`, such so that when the async sequence retruned by the closure
throws, it throws the same thing as the base sequence. Therefore, the
`Failure` type is defined to be `Base.Failure`.

This property isn't enforced at the type level, but instead in the
`AsyncSequence.flatMap` signatures: we replace the one signature that
returned `AsyncFlatMapSequence` with three overloads that differ only
in their generic requirements, adding:
1. `where SegmentOfResult.Failure == Failure`
2. `where SegmentOfResult.Failure == Never`
3. `where SegmentOfResult.Failure == Never, Failure == Never` (a
tiebreaker between the two above)

For cases where `SegmentOfResult.Failure` is neither `Never` nor
`Failure`, overloading will choose the `flatMap` function that returns
an `AsyncThrowingFlatMapSequence`. This can mean that existing code
will choose a different overload and get a different type, but other
than the type identity changing, the resulting sequence will behave
the same way.

Author: DougGregor

stdlib/public/Concurrency/AsyncFlatMapSequence.swift

@@ -14,6 +14,38 @@ import Swift
 
 @available(SwiftStdlib 5.1, *)
 extension AsyncSequence {
+  /// Creates an asynchronous sequence that concatenates the results of calling
+  /// the given transformation with each element of this sequence.
+  ///
+  /// Use this method to receive a single-level asynchronous sequence when your
+  /// transformation produces an asynchronous sequence for each element.
+  ///
+  /// In this example, an asynchronous sequence called `Counter` produces `Int`
+  /// values from `1` to `5`. The transforming closure takes the received `Int`
+  /// and returns a new `Counter` that counts that high. For example, when the
+  /// transform receives `3` from the base sequence, it creates a new `Counter`
+  /// that produces the values `1`, `2`, and `3`. The `flatMap(_:)` method
+  /// "flattens" the resulting sequence-of-sequences into a single
+  /// `AsyncSequence`.
+  ///
+  ///     let stream = Counter(howHigh: 5)
+  ///         .flatMap { Counter(howHigh: $0) }
+  ///     for await number in stream {
+  ///         print(number, terminator: " ")
+  ///     }
+  ///     // Prints "1 1 2 1 2 3 1 2 3 4 1 2 3 4 5 "
+  ///
+  /// - Parameter transform: A mapping closure. `transform` accepts an element
+  ///   of this sequence as its parameter and returns an `AsyncSequence`.
+  /// - Returns: A single, flattened asynchronous sequence that contains all
+  ///   elements in all the asynchronous sequences produced by `transform`.
+  @usableFromInline
+  __consuming func flatMap<SegmentOfResult: AsyncSequence>(
+    _ transform: @Sendable @escaping (Element) async -> SegmentOfResult
+  ) -> AsyncFlatMapSequence<Self, SegmentOfResult> {
+    return AsyncFlatMapSequence(self, transform: transform)
+  }
+
   /// Creates an asynchronous sequence that concatenates the results of calling
   /// the given transformation with each element of this sequence.
   ///
@@ -40,12 +72,88 @@ extension AsyncSequence {
   /// - Returns: A single, flattened asynchronous sequence that contains all
   ///   elements in all the asynchronous sequences produced by `transform`.
   @preconcurrency 
+  @_alwaysEmitIntoClient
   @inlinable
   public __consuming func flatMap<SegmentOfResult: AsyncSequence>(
     _ transform: @Sendable @escaping (Element) async -> SegmentOfResult
-  ) -> AsyncFlatMapSequence<Self, SegmentOfResult> {
+  ) -> AsyncFlatMapSequence<Self, SegmentOfResult>
+    where SegmentOfResult.Failure == Failure
+  {
+    return AsyncFlatMapSequence(self, transform: transform)
+  }
+
+  /// Creates an asynchronous sequence that concatenates the results of calling
+  /// the given transformation with each element of this sequence.
+  ///
+  /// Use this method to receive a single-level asynchronous sequence when your
+  /// transformation produces an asynchronous sequence for each element.
+  ///
+  /// In this example, an asynchronous sequence called `Counter` produces `Int`
+  /// values from `1` to `5`. The transforming closure takes the received `Int`
+  /// and returns a new `Counter` that counts that high. For example, when the
+  /// transform receives `3` from the base sequence, it creates a new `Counter`
+  /// that produces the values `1`, `2`, and `3`. The `flatMap(_:)` method
+  /// "flattens" the resulting sequence-of-sequences into a single
+  /// `AsyncSequence`.
+  ///
+  ///     let stream = Counter(howHigh: 5)
+  ///         .flatMap { Counter(howHigh: $0) }
+  ///     for await number in stream {
+  ///         print(number, terminator: " ")
+  ///     }
+  ///     // Prints "1 1 2 1 2 3 1 2 3 4 1 2 3 4 5 "
+  ///
+  /// - Parameter transform: A mapping closure. `transform` accepts an element
+  ///   of this sequence as its parameter and returns an `AsyncSequence`.
+  /// - Returns: A single, flattened asynchronous sequence that contains all
+  ///   elements in all the asynchronous sequences produced by `transform`.
+  @preconcurrency 
+  @_alwaysEmitIntoClient
+  @inlinable
+  public __consuming func flatMap<SegmentOfResult: AsyncSequence>(
+    _ transform: @Sendable @escaping (Element) async -> SegmentOfResult
+  ) -> AsyncFlatMapSequence<Self, SegmentOfResult>
+    where SegmentOfResult.Failure == Never
+  {
     return AsyncFlatMapSequence(self, transform: transform)
   }
+
+  /// Creates an asynchronous sequence that concatenates the results of calling
+  /// the given transformation with each element of this sequence.
+  ///
+  /// Use this method to receive a single-level asynchronous sequence when your
+  /// transformation produces an asynchronous sequence for each element.
+  ///
+  /// In this example, an asynchronous sequence called `Counter` produces `Int`
+  /// values from `1` to `5`. The transforming closure takes the received `Int`
+  /// and returns a new `Counter` that counts that high. For example, when the
+  /// transform receives `3` from the base sequence, it creates a new `Counter`
+  /// that produces the values `1`, `2`, and `3`. The `flatMap(_:)` method
+  /// "flattens" the resulting sequence-of-sequences into a single
+  /// `AsyncSequence`.
+  ///
+  ///     let stream = Counter(howHigh: 5)
+  ///         .flatMap { Counter(howHigh: $0) }
+  ///     for await number in stream {
+  ///         print(number, terminator: " ")
+  ///     }
+  ///     // Prints "1 1 2 1 2 3 1 2 3 4 1 2 3 4 5 "
+  ///
+  /// - Parameter transform: A mapping closure. `transform` accepts an element
+  ///   of this sequence as its parameter and returns an `AsyncSequence`.
+  /// - Returns: A single, flattened asynchronous sequence that contains all
+  ///   elements in all the asynchronous sequences produced by `transform`.
+  @preconcurrency 
+  @_alwaysEmitIntoClient
+  @inlinable
+  public __consuming func flatMap<SegmentOfResult: AsyncSequence>(
+    _ transform: @Sendable @escaping (Element) async -> SegmentOfResult
+  ) -> AsyncFlatMapSequence<Self, SegmentOfResult>
+  where SegmentOfResult.Failure == Never, Failure == Never
+  {
+    return AsyncFlatMapSequence(self, transform: transform)
+  }
+
 }
 
 /// An asynchronous sequence that concatenates the results of calling a given
@@ -75,6 +183,14 @@ extension AsyncFlatMapSequence: 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 the type of error in the base asynchronous
+  /// sequence. By construction, the sequence produced by the `transform`
+  /// closure must either produce this type of error or not produce errors
+  /// at all.
+  @available(SwiftStdlib 5.11, *)
+  public typealias Failure = Base.Failure
   /// The type of iterator that produces elements of the sequence.
   public typealias AsyncIterator = Iterator
 
@@ -148,6 +264,58 @@ extension AsyncFlatMapSequence: AsyncSequence {
       }
       return nil
     }
+
+    /// Produces the next element in the flat 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, takes the
+    /// resulting asynchronous sequence, and creates an asynchronous iterator
+    /// from it.  `nextElement()` then consumes values from this iterator until
+    /// it terminates.  At this point, `nextElement()` is ready to receive the
+    /// next value from the base sequence.
+    @available(SwiftStdlib 5.11, *)
+    @inlinable
+    public mutating func nextElement() async throws(Failure) -> SegmentOfResult.Element? {
+      while !finished {
+        if var iterator = currentIterator {
+          do throws(any Error) {
+            let optElement = try await iterator.nextElement()
+            guard let element = optElement else {
+              currentIterator = nil
+              continue
+            }
+            // restore the iterator since we just mutated it with next
+            currentIterator = iterator
+            return element
+          } catch {
+            finished = true
+            throw error as! Failure
+          }
+        } else {
+          let optItem = try await baseIterator.nextElement()
+          guard let item = optItem else {
+            finished = true
+            return nil
+          }
+          do throws(any Error) { 
+            let segment = await transform(item)
+            var iterator = segment.makeAsyncIterator()
+            let optElement = try await iterator.nextElement()  
+            guard let element = optElement else {
+              currentIterator = nil
+              continue
+            }
+            currentIterator = iterator
+            return element
+          } catch {
+            finished = true
+            throw error as! Failure
+          }
+        }
+      }
+      return nil
+    }
   }
 
   @inlinable