Swift 中 map & flatMap & compactMap 的区别

2019-06-13
Swift

Swift 中提供了很多的高阶函数，而其中map，flatMap 与 compactMap 又有许多的相似之处，经常对他们使用的时机有些模糊，今天我们就从标准库实现的角度来看下他们各自的区别以及适用的情况。

Overview

关于它们的定义：

1
2
3
4
5
6
// Sequence
public func map<T>(_ transform: (Element) throws -> T) rethrows -> [T]

public func flatMap<SegmentOfResult : Sequence>(_ transform: (Element) throws -> SegmentOfResult) rethrows -> [SegmentOfResult.Element]

public func compactMap<ElementOfResult>(_ transform: (Element) throws -> ElementOfResult?) rethrows -> [ElementOfResult]

只是从上面的函数签名来看这三个函数，它们虽然都接受一个 transform 的闭包作为参数，且闭包捕获的参数都为 Sequence 中的元素类型，但是闭包所返回的类型可以与 Sequence 元素的类型不同。
map 中的闭包返回类型为 T，而 map 函数的返回类型为 [T]。

flatMap 中的闭包的参数同样是 Sequence 中的元素类型，但其返回类型为 SegmentOfResult。在函数体的范型定义中，SegmentOfResult 的类型其实就是是 Sequence。 而flatMap 函数返回的类型是 SegmentOfResult.Element 的数组。从函数的返回值来看，与 map 的区别在于 flatMap 会将 Sequence 中的元素进行”降温”，返回的类型会是 Sequence 中元素类型的数组，而 map 返回的这是闭包返回类型T的数组。

compactMap 闭包参数的返回类型为 ElementOfResult?, 而函数返回类型为 [ElementOfResult]。因此，compactMap 会将闭包返回的 Optional 的类型解包并返回解包后的数组。

下面分别从源码角度再看下标准库的具体实现方式：

map

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
public func map<T>(
    _ transform: (Element) throws -> T
  ) rethrows -> [T] {
    // underestimatedCount 返回值的大小会略小于或等于当前序列元素的数量
    let initialCapacity = underestimatedCount
    var result = ContiguousArray<T>()
    // 保留当前序列预估容量的存储空间
    result.reserveCapacity(initialCapacity)

    var iterator = self.makeIterator()

    // Add elements up to the initial capacity without checking for regrowth.
    // initialCapacity 会略小于当前序列的数量，因此这一段的遍历可以使用强制解包`interatpr.next()`
    for _ in 0..<initialCapacity {
      result.append(try transform(iterator.next()!))
    }
    // Add remaining elements, if any.
    while let element = iterator.next() {
      result.append(try transform(element))
    }
    return Array(result)
  }

在上面 map 的实现中，主要做了以下几件事：

  • 首先根据序列中的 underestimatedCount 和变换的目标类型 T，初始化一个 result 的数组来存放结果。
  • 通过 iterator 遍历序列中的每一个元素，并调用 transform 闭包对元素进行变换。
  • 将把变换后的结果保存结果数组中。

因此，经过 map 变换的 Sequence 就不再是一个简单的序列了，而是一个数组。我们只能对有限序列使用map进行变换。

ContiguousArray 和名字所暗示的不同，它其实是 Swift 中最简单的数组类型。相比标准的数组，它可以有更好的性能表现，而即便没有，也至少可以提供与 Array 相同性能水平的表现。同时也暴露出相同的接口。

flatMap

1
2
3
4
5
6
7
8
9
public func flatMap<SegmentOfResult : Sequence>(
    _ transform: (Element) throws -> SegmentOfResult
  ) rethrows -> [SegmentOfResult.Element] {
    var result: [SegmentOfResult.Element] = []
    for element in self {
      result.append(contentsOf: try transform(element))
    }
    return result
  }

对于 flatMap 的实现，我们可以看出，它做了以下几件事情：

  • 初始化一个名为 result 的新数组，用于存放结果。
  • 遍历自己的元素，对于每个元素，调用闭包的转换函数 transform 进行转换。
  • 将转换的结果，使用 appendContentsOf 方法，将结果放入 result 数组中。

而这个 appendContentsOf 方法，即是把数组中的元素取出来，放入新数组。

compactMap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
public func compactMap<ElementOfResult>(
    _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    return try _compactMap(transform)
  }

public func _compactMap<ElementOfResult>(
    _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    var result: [ElementOfResult] = []
    for element in self {
      if let newElement = try transform(element) {
        result.append(newElement)
      }
    }
    return result
  }

而compactMap 的实现则做了以下事情：

  • 构造一个名为 result 的新数组，用于存放结果。
  • 遍历自己的元素，对于每个元素，调用闭包的闭包参数 transform，进行转换。
  • 将转换的结果进行解包，如果有值则使用 append 方法，将结果放入 result 数组中。

所以，该 compactMap 函数可以过滤闭包执行结果为 nil 的情况，仅收集那些转换后非空的结果。

如何选择他们使用的时机呢？

When to use compactMap

当转换闭包返回可选值并且你期望得到的结果为非可选值的序列时，使用 compactMap。

compactMap 与 map 的区别参考下面的例子：

1
2
3
4
5
6
7
let scores = ["1", "2", "three", "four", "5"]

let mapped: [Int?] = scores.map { str in Int(str) }
// [1, 2, nil, nil, 5] - Two nil values as "three" and "four" are strings.

let compactMapped: [Int] = scores.compactMap { str in Int(str) }
// [1, 2, 5] - The nil values for "three" and "four" are filtered out.

根据这些不同点，可以归纳出三个函数适用的不同情况，当闭包返回的结果需要是序列类型且期望返回的结果是一维数组时，使用 flatMap, 而需要过滤 transform 结果的可选值时，使用 compactMap。

如：

1
2
3
4
5
6
7
8
9
10
11
12
13
let arr = [[1, 2, 3], [4, 5]]

let result = arr.map { $0 }
// [[1, 2, 3], [4, 5]]

let result = arr.flatMap { $0 }
// [1, 2, 3, 4, 5]


let arr = [1, 2, 3, nil, nil, 4, 5]

let result = arr.compactMap { $0 }
// [1, 2, 3, 4, 5]

When to use flatMap

当对于序列中元素，转换闭包返回的是序列或者集合时，而你期望得到的结果是一维数组时，使用 flatMap。

1
2
3
4
5
6
7
8
let scoresByName = ["Henk": [0, 5, 8], "John": [2, 5, 8]]

let mapped = scoresByName.map { $0.value }
// [[0, 5, 8], [2, 5, 8]] - An array of arrays
print(mapped)

let flatMapped = scoresByName.flatMap { $0.value }
// [0, 5, 8, 2, 5, 8] - flattened to only one array

其实 s.flatMap(transform) 的结果等同于 Array(s.map(transform).joined())。

compactMap vs flatMap

当在序列元素上使用转换闭包，且返回值为Optional 类型时，使用 compactMap, 其他情况下 map 与 flatMap 都可以得到你所想要的结果。