Merge pull request #2461 from russbishop/fork/master

Author: swift-ci

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<ElementOfResult>(
+    _ transform: (Elements.Iterator.Element) -> ElementOfResult?
+  ) -> LazyMapSequence<
+    LazyFilterSequence<
+      LazyMapSequence<Elements, ElementOfResult?>>,
+    ElementOfResult
+  > {
+    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<ElementOfResult>(
+    _ transform: (Elements.Iterator.Element) -> ElementOfResult?
+  ) -> LazyMapCollection<
+    LazyFilterCollection<
+      LazyMapCollection<Elements, ElementOfResult?>>,
+    ElementOfResult
+  > {
+    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<ElementOfResult>(
+    _ transform: (Elements.Iterator.Element) -> ElementOfResult?
+  ) -> LazyMapBidirectionalCollection<
+    LazyFilterBidirectionalCollection<
+      LazyMapBidirectionalCollection<Elements, ElementOfResult?>>,
+    ElementOfResult
+  > {
+    return self.map(transform).filter { $0 != nil }.map { $0! }
+  }
 }

stdlib/public/core/SequenceAlgorithms.swift.gyb

@@ -657,10 +657,10 @@ extension Sequence {
   /// - Complexity: O(*m* + *n*), where *m* is the length of this sequence
   ///   and *n* is the length of the result.
   @warn_unused_result
-  public func flatMap<T>(
-    _ transform: @noescape (${GElement}) throws -> T?
-  ) rethrows -> [T] {
-    var result: [T] = []
+  public func flatMap<ElementOfResult>(
+    _ transform: @noescape (${GElement}) throws -> ElementOfResult?
+  ) rethrows -> [ElementOfResult] {
+    var result: [ElementOfResult] = []
     for element in self {
       if let newElement = try transform(element) {
         result.append(newElement)

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