[stdlib] added flatMap

fixes rdar://problem/19338087

Swift SVN r25378

stdlib/core/Arrays.swift.gyb

@@ -608,6 +608,16 @@ extension ${Self} : _ArrayType {
     return ${Self}<U>(lazy(self).map(transform))
   }
 
+  /// Return ${a_Self} containing the results of calling
+  /// `transform(x)` on each element `x` of `self` and flattening the result.
+  public func flatMap<U>(@noescape transform: (T) -> ${Self}<U>) -> ${Self}<U> {
+    var result: ${Self}<U> = []
+    for element in self {
+      result.extend(transform(element))
+    }
+    return result
+  }
+
   /// A ${Self} containing the elements of `self` in reverse order
   public func reverse() -> ${Self} {
     return ${Self}(lazy(self).reverse())

stdlib/core/ImplicitlyUnwrappedOptional.swift

@@ -54,6 +54,22 @@ public enum ImplicitlyUnwrappedOptional<T>
     }
   }
 
+  /// Returns `f(self)!` iff `self` and `f(self)` are not nil.
+  public func flatMap<U>(@noescape f: (T) -> ImplicitlyUnwrappedOptional<U>)
+    -> ImplicitlyUnwrappedOptional<U> {
+    switch self {
+    case .Some(let y):
+      switch f(y) {
+      case .Some(let z):
+        return .Some(z)
+      case .None:
+        return .None
+      }
+    case .None:
+      return .None
+    }
+  }
+
   /// Returns a mirror that reflects `self`.
   public func getMirror() -> MirrorType {
     // FIXME: This should probably use _OptionalMirror in both cases.

stdlib/core/Map.swift.gyb

@@ -73,6 +73,18 @@ public func map<S : SequenceType, T>(
   return lazy(source).map(transform).array
 }
 
+/// Return an `Array` containing the results of mapping `transform`
+/// over `source` and flattening the result.
+public func flatMap<S : SequenceType, T>(
+  source: S, @noescape transform: (S.Generator.Element) -> [T]
+) -> [T] {
+  var result: [T] = []
+  for element in source {
+    result.extend(transform(element))
+  }
+  return result
+}
+
 //===--- Collections ------------------------------------------------------===//
 
 /// A `CollectionType` whose elements consist of those in a `Base`
@@ -129,6 +141,18 @@ public func map<C : CollectionType, T>(
   return lazy(source).map(transform).array
 }
 
+/// Return an `Array` containing the results of mapping `transform`
+/// over `source` and flattening the result.
+public func flatMap<C : CollectionType, T>(
+  source: C, transform: (C.Generator.Element) -> [T]
+) -> [T] {
+  var result: [T] = []
+  for elements in map(source, transform) {
+    result.extend(elements)
+  }
+  return result
+}
+
 //===--- Support for lazy(s) ----------------------------------------------===//
 
 % traversals = ('Forward', 'Bidirectional', 'RandomAccess')

stdlib/core/Optional.swift

@@ -34,6 +34,21 @@ public enum Optional<T> : Reflectable, NilLiteralConvertible {
     }
   }
 
+  /// Returns `f(self)!` iff `self` and `f(self)` are not nil.
+  public func flatMap<U>(@noescape f: (T)->U?) -> U? {
+    switch self {
+    case .Some(let y):
+      switch f(y) {
+      case .Some(let z):
+        return .Some(z)
+      case .None:
+        return .None
+      }
+    case .None:
+      return .None
+    }
+  }
+
   /// Returns a mirror that reflects `self`.
   public func getMirror() -> MirrorType {
     return _OptionalMirror(self)
@@ -78,6 +93,22 @@ public func map<T, U>(x: T?, @noescape f: (T)->U) -> U? {
   }
 }
 
+
+/// Returns `f(self)!` iff `self` and `f(self)` are not nil.
+public func flatMap<T, U>(x: T?, @noescape f: (T)->U?) -> U? {
+  switch x {
+  case .Some(let y):
+    switch f(y) {
+    case .Some(let z):
+      return .Some(z)
+    case .None:
+      return .None
+    }
+  case .None:
+    return .None
+  }
+}
+
 // Intrinsics for use by language features.
 @transparent
 public // COMPILER_INTRINSIC

test/1_stdlib/Algorithm.swift

@@ -319,6 +319,45 @@ Algorithm.test("map/CollectionType") {
   }
 }
 
+// FIXME(19789546): move to FlatMap test file
+Algorithm.test("flatMap/SequenceType") {
+  if true {
+    let s = DrainableSequence<Int>([])
+    var result = flatMap(s) {
+      (x: Int) -> [Int16] in
+      expectUnreachable()
+      return [42]
+    }
+    expectType([Int16].self, &result)
+    expectEqual([], result)
+    expectEqual([], Array(s))
+  }
+  if true {
+    let s = DrainableSequence([ 0, 30, 10, 90 ])
+    let result = flatMap(s) { [$0 + 1] }
+    expectEqual([ 1, 31, 11, 91 ], result)
+    expectEqual([], Array(s))
+  }
+}
+
+Algorithm.test("flatMap/CollectionType") {
+  if true {
+    let c = MinimalForwardCollection<Int>([])
+    var result = flatMap(c) {
+      (x: Int) -> [Int16] in
+      expectUnreachable()
+      return [42]
+    }
+    expectType([Int16].self, &result)
+    expectEqual([], result)
+  }
+  if true {
+    let c = MinimalForwardCollection([ 0, 30, 10, 90 ])
+    let result = flatMap(c) { [$0 + 1] }
+    expectEqual([ 1, 31, 11, 91 ], result)
+  }
+}
+
 Algorithm.test("sorted/strings") {
   expectEqual(
     [ "Banana", "apple", "cherry" ],

test/1_stdlib/ImplicitlyUnwrappedOptional.swift

@@ -36,3 +36,24 @@ if c === nil {
   println("x is not nil!")
 }
 // CHECK: an empty class optional should equal nil
+
+import StdlibUnittest
+import Swift
+
+var ImplicitlyUnwrappedOptionalTests = TestSuite("ImplicitlyUnwrappedOptional")
+
+let half : Int -> Int! =
+  { if $0 % 2 == 0 { return $0 / 2 } else { return .None } }
+
+ImplicitlyUnwrappedOptionalTests.test("flatMap") {
+  // FIXME(19798684): can't call map or flatMap on ImplicitlyUnwrappedOptional
+  // expectTrue(ImplicitlyUnwrappedOptional<Int>.None.flatMap(half) ==
+  //   ImplicitlyUnwrappedOptional<Int>.None)
+  // expectTrue(half(4) == ImplicitlyUnwrappedOptional<Int>.Some(2))
+  // expectTrue(half(4).flatMap(half) == ImplicitlyUnwrappedOptional<Int>.Some(1))
+  // expectTrue(half(4).flatMap(half).flatMap(half) ==
+  //   ImplicitlyUnwrappedOptional<Int>.None)
+}
+
+runAllTests()
+

test/1_stdlib/Optional.swift

@@ -203,3 +203,21 @@ debugPrintln(c ?? nextCounter2())      // CHECK-NEXT: Optional(4)
 debugPrintln(d ?? nextCounter2())      // CHECK-NEXT: Optional(456)
 debugPrintln(e ?? d ?? nextCounter2()) // CHECK-NEXT: Optional(456)
 debugPrintln(f ?? nextCounter2())      // CHECK-NEXT: Optional(5)
+
+import StdlibUnittest
+import Swift
+
+var OptionalTests = TestSuite("Optional")
+
+let half : Int -> Int? =
+  { if $0 % 2 == 0 { return $0 / 2 } else { return .None } }
+
+OptionalTests.test("flatMap") {
+  // FIXME: type inference stops expectEqual from working
+  expectTrue((.None as Int?).flatMap(half) == .None)
+  expectTrue(half(4) == .Some(2))
+  expectTrue(half(4).flatMap(half) == .Some(1))
+  expectTrue(half(4).flatMap(half).flatMap(half) == .None)
+}
+
+runAllTests()

validation-test/stdlib/ArrayNew.swift.gyb

@@ -188,6 +188,17 @@ ArrayTestSuite.test("${array_type}/map") {
   }
 }
 
+ArrayTestSuite.test("${array_type}/flatMap") {
+  let enumerate : Int -> ${array_type}<Int> =
+    { return ${array_type}(1..<($0 + 1)) }
+  expectEqual(${array_type}([]), ${array_type}().flatMap(enumerate))
+  expectEqual(${array_type}([1]), ${array_type}([1]).flatMap(enumerate))
+  expectEqual(${array_type}([1, 1, 2]),
+    ${array_type}([1, 2]).flatMap(enumerate))
+  expectEqual(${array_type}([1, 1, 1, 2]),
+    ${array_type}([1, 2]).flatMap(enumerate).flatMap(enumerate))
+}
+
 //===---
 // Check that generators traverse a snapshot of the collection.
 //===---