[stdlib] SR-361: Implement SE008: Add Lazy flatMap for Seq of Optionals

stdlib/public/core/FlatMap.swift

@@ -24,6 +24,25 @@ extension LazySequenceProtocol {
     FlattenSequence<LazyMapSequence<Elements, SegmentOfResult>>> {
     return self.map(transform).flatten()
   }
+
+  /// Returns a `LazyMapSequence` containing the concatenated non-nil
+  /// results of mapping transform over this `Sequence`.
+  ///
+  /// Use this method to receive only nonoptional values when your
+  /// transformation produces an optional value.
+  ///
+  /// - Parameter transform: A closure that accepts an element of this
+  /// sequence as its argument and returns an optional value.
+  @warn_unused_result
+  public func flatMap<U>(
+    _ transform: (Self.Elements.Iterator.Element) -> U?
+  ) -> LazyMapSequence<
+    LazyFilterSequence<
+      LazyMapSequence<Self.Elements, U?>>,
+    U
+  > {
+    return self.map(transform).filter { $0 != nil }.map { $0! }
+  }
 }
 
 extension LazyCollectionProtocol {
@@ -42,6 +61,25 @@ extension LazyCollectionProtocol {
   > {
     return self.map(transform).flatten()
   }
+
+  /// Returns a `LazyMapCollection` containing the concatenated non-nil
+  /// results of mapping transform over this collection.
+  ///
+  /// Use this method to receive only nonoptional values when your
+  /// transformation produces an optional value.
+  ///
+  /// - Parameter transform: A closure that accepts an element of this
+  /// collection as its argument and returns an optional value.
+  @warn_unused_result
+  public func flatMap<U>(
+    _ transform: (Self.Elements.Iterator.Element) -> U?
+  ) -> LazyMapCollection<
+    LazyFilterCollection<
+      LazyMapCollection<Self.Elements, U?>>,
+    U
+  > {
+    return self.map(transform).filter { $0 != nil }.map { $0! }
+  }
 }
 
 extension LazyCollectionProtocol
@@ -67,4 +105,22 @@ extension LazyCollectionProtocol
   >> {
     return self.map(transform).flatten()
   }
+
+  /// Returns a `LazyMapBidirectionalCollection` containing the concatenated non-nil
+  /// results of mapping transform over this collection.
+  ///
+  /// Use this method to receive only nonoptional values when your
+  /// transformation produces an optional value.
+  ///
+  /// - Parameter transform: A closure that accepts an element of this
+  /// collection as its argument and returns an optional value.
+  public func flatMap<U>(
+    _ transform: (Self.Elements.Iterator.Element) -> U?
+  ) -> LazyMapBidirectionalCollection<
+    LazyFilterBidirectionalCollection<
+      LazyMapBidirectionalCollection<Self.Elements, U?>>,
+    U
+  > {
+    return self.map(transform).filter { $0 != nil }.map { $0! }
+  }
 }

test/1_stdlib/Inputs/flatMap.gyb

@@ -64,5 +64,29 @@ ${Test}.test("flatMap/${Kind}/Lazy") {
       + "once per element")
     % end
   }
+
+  for test in flatMapToOptionalTests {
+    let s = ${Minimal}<OpaqueValue<Int>>(elements:
+      test.sequence.map(OpaqueValue.init))
+    let closureLifetimeTracker = LifetimeTracked(0)
+    var timesClosureWasCalled = 0
+
+    var result = s.lazy.flatMap {
+      (element: OpaqueValue<Int>) -> OpaqueValue<Int32>? in
+      _blackHole(closureLifetimeTracker)
+      timesClosureWasCalled += 1
+      return test.transform(element.value).map(OpaqueValue.init)
+    }
+    expectEqual(0, timesClosureWasCalled, "unexpected eagerness")
+
+    let expected = test.expected.map(OpaqueValue.init)
+    expectEqualSequence(
+      expected, result,
+      stackTrace: SourceLocStack().with(test.loc)
+    ) { $0.value == $1.value }
+    expectEqual(
+      test.sequence.count, timesClosureWasCalled,
+      "iterating lazy flatMap() should call closure once per element")
+  }
 }
 % end