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

stdlib/public/core/Filter.swift.gyb

@@ -165,7 +165,7 @@ public func > <Base : Collection>(
 // collapsed into one `LazyFilterCollection` using conditional conformances.
 // Maybe even combined with `LazyFilterSequence`.
 
-% for Traversal in ['Forward', 'Bidirectional']:
+% for Traversal in ['Forward', 'Bidirectional', 'RandomAccess']:
 %   Self = "LazyFilter" + collectionForTraversal(Traversal)
 %   Slice = sliceTypeName(traversal=Traversal, mutable=False, rangeReplaceable=False)
 
@@ -244,7 +244,7 @@ public struct ${Self}<
     i = LazyFilterIndex(base: index)
   }
 
-%   if Traversal == 'Bidirectional':
+%   if Traversal == 'Bidirectional' or Traversal == 'RandomAccess':
   @warn_unused_result
   public func index(before i: Index) -> Index {
     var i = i
@@ -315,7 +315,7 @@ extension LazySequenceProtocol {
   }
 }
 
-% for Traversal in ['Forward', 'Bidirectional']:
+% for Traversal in ['Forward', 'Bidirectional', 'RandomAccess']:
 
 extension LazyCollectionProtocol
   where

stdlib/public/core/FlatMap.swift

@@ -24,6 +24,26 @@ 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`.  The elements of
+  /// the result are computed lazily, each time they are read.
+  ///
+  /// 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 +62,26 @@ extension LazyCollectionProtocol {
   > {
     return self.map(transform).flatten()
   }
+
+  /// Returns a `LazyMapCollection` containing the concatenated non-nil
+  /// results of mapping transform over this collection.  The elements of
+  /// the result are computed lazily, each time they are read.
+  ///
+  /// 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 +107,48 @@ extension LazyCollectionProtocol
   >> {
     return self.map(transform).flatten()
   }
+
+  /// Returns a `LazyMapBidirectionalCollection` containing the concatenated non-nil
+  /// results of mapping transform over this collection.  The elements of
+  /// the result are computed lazily, each time they are read.
+  ///
+  /// 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! }
+  }
+}
+
+extension LazyCollectionProtocol
+  where
+  Self : RandomAccessCollection,
+  Elements : RandomAccessCollection
+{
+  /// Returns a `LazyMapRandomAccessCollection` containing the concatenated non-nil
+  /// results of mapping transform over this collection.  The elements of
+  /// the result are computed lazily, each time they are read.
+  ///
+  /// 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?
+  ) -> LazyMapRandomAccessCollection<
+    LazyFilterRandomAccessCollection<
+      LazyMapRandomAccessCollection<Self.Elements, U?>>,
+    U
+  > {
+    return self.map(transform).filter { $0 != nil }.map { $0! }
+  }
 }

validation-test/stdlib/Lazy.swift.gyb

@@ -1161,6 +1161,20 @@ tests.test("LazyFilterBidirectionalCollection/AssociatedTypes") {
     indicesType: DefaultBidirectionalIndices<Subject>.self)
 }
 
+tests.test("LazyFilterRandomAccessCollection/AssociatedTypes") {
+  typealias Base = MinimalRandomAccessCollection<OpaqueValue<Int>>
+  typealias Subject = LazyFilterRandomAccessCollection<Base>
+  expectRandomAccessCollectionAssociatedTypes(
+    collectionType: Subject.self,
+    iteratorType: LazyFilterIterator<Base.Iterator>.self,
+    // FIXME(ABI): SubSequence should be `LazyFilterRandomAccessCollection<Base.Slice>`.
+    subSequenceType: RandomAccessSlice<Subject>.self,
+    indexType: LazyFilterIndex<Base>.self,
+    indexDistanceType: Base.IndexDistance.self,
+    indicesType: DefaultRandomAccessIndices<Subject>.self)
+}
+
+
 tests.test("lazy.filter/TypeInference") {
   let baseArray: [OpaqueValue<Int>] = (0..<10).map(OpaqueValue.init)
   do {
@@ -1190,7 +1204,7 @@ tests.test("lazy.filter/TypeInference") {
     var filtered = MinimalRandomAccessCollection(elements: baseArray)
       .lazy.filter { _ in true }
     expectType(
-      LazyFilterBidirectionalCollection<
+      LazyFilterRandomAccessCollection<
         MinimalRandomAccessCollection<OpaqueValue<Int>>
       >.self,
       &filtered)
@@ -1283,5 +1297,22 @@ do {
   }
 }
 
-runAllTests()
+% for Traversal in TRAVERSALS:
+%		TraversalCollection = collectionForTraversal(Traversal)
+tests.test("${TraversalCollection}/lazy/flatMap") {
+	let sample = (0..<100)
+	let expected = sample.filter { $0 % 3 != 0 }.map { String($0) }
+	let base = Minimal${TraversalCollection}(elements: sample).lazy
+
+	var calls = 0
+	let mapped = sample.lazy.flatMap { (x) -> String? in
+		calls += 1
+		return x % 3 != 0 ? String(x) : nil
+	}
+	expectEqual(0, calls, "unexpected eagerness in \(base.dynamicType).flatMap")
+
+	checkSequence(expected, mapped, resiliencyChecks: .none)
+}
+%end
 
+runAllTests()