Swift 源码解读 - Map.swift

map 、 flatMap

map 和 flatMap 是 Swift 中常用的高阶函数，下面将从源码的角度来解析使用，需要注意的是 map 和 flatMap 都是 Sequence 协议的扩展方法，并不是 Map.swift 里面的 map 方法，这两个方法的源码解读全为了下文作铺垫。

map 方法

map 接受一个闭包作为参数，作用于数组中的每个元素，然后将这些变换后的元素组成一个新的数组并返回：

let array = [1, 2, 3].map { $0 * 2 } // [2, 4, 6]

我们来看下 Swift 中这个方法的源码：

// Sequence.swift
@_inlineable
public func map<T>(
    _ transform: (Element) throws -> T
    ) rethrows -> [T] {
    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.
    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 是 Sequence 协议的一个扩展方法，源码也很简单，直接对每个元素 transform，然后将 transform 的元素组成一个新的数组。

flatMap

关于 flatMap 方法，先从源码来分析，暂不淡使用。flatMap 方法在 SequenceAlgorithms.swift 中，它有两个重载：

• 重载 1

  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 方法和 map 方法非常相似的，区别就是 append(contentsOf: ) 和 append()，所以它具有降维的作用：

  [[1,2,3], [4, 5]].flatMap { $0.map { $0 * 2} } // [2, 4, 6, 8, 10]
  [1, 2, 3].flatMap({ $0 * 2 }) // [2, 4, 6]
  

• 重载 2

  @_inlineable
  public func flatMap<ElementOfResult>(
  _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    return try _flatMap(transform)
  }
  
  @_inlineable // FIXME(sil-serialize-all)
  @inline(__always)
  public func _flatMap<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
  }
  

  这个方法的关键代码是：

  if let newElement = try transform(element) {
   result.append(newElement)
  }
  

  如果数组中存在可选值或者 nil 时，系统就会调用这个重载，会把 optional 强制解包并把 nil 过滤掉：

  [1, 2, nil, Optional(3)].flatMap { $0 }  // [1, 2, 3]
  

通过阅读 map 和 flatMap 的源码，我们很易容就学会了这两个高阶函数的使用，铺垫也说了，下面来到今天主要想说的内容。

Map.swift 源码解读

这里我们从上向下读下 Map.swift 的源码，先看 LazyMapSequence 源码：

public struct LazyMapSequence<Base : Sequence, Element>
  : LazySequenceProtocol {

  public typealias Elements = LazyMapSequence

  /// - Complexity: O(1).
  @_inlineable
  public func makeIterator() -> LazyMapIterator<Base.Iterator, Element> {
    return LazyMapIterator(_base: _base.makeIterator(), _transform: _transform)
  }

  /// - Complexity: O(*n*)
  @_inlineable
  public var underestimatedCount: Int {
    return _base.underestimatedCount
  }

  @_inlineable
  @_versioned
  internal init(_base: Base, transform: @escaping (Base.Element) -> Element) {
    self._base = _base
    self._transform = transform
  }

  @_versioned
  internal var _base: Base
  @_versioned
  internal let _transform: (Base.Element) -> Element
}

LazyMapSequence 实现了 LazySequenceProtocol，而 LazySequenceProtocol 则是实现了 Sequence，同样 LazyMapSequence 需要实现 Sequence 协议，所以必需实现 makeIterator 方法：

public func makeIterator() -> LazyMapIterator<Base.Iterator, Element> {
  return LazyMapIterator(_base: _base.makeIterator(), _transform: _transform)
}

关键点是这里使用了新定义的 LazyMapIterator 迭代器，我们跳转到 LazyMapIterator 源码处查看：

public struct LazyMapIterator<
  Base : IteratorProtocol, Element
> : IteratorProtocol, Sequence {
  /// Advances to the next element and returns it, or `nil` if no next element
  /// exists.
  ///
  /// Once `nil` has been returned, all subsequent calls return `nil`.
  ///
  /// - Precondition: `next()` has not been applied to a copy of `self`
  ///   since the copy was made.
  @_inlineable
  public mutating func next() -> Element? {
    return _base.next().map(_transform)
  }

  @_inlineable
  public var base: Base { return _base }

  @_versioned
  internal var _base: Base
  @_versioned
  internal let _transform: (Base.Element) -> Element

  @_inlineable
  @_versioned
  internal init(_base: Base, _transform: @escaping (Base.Element) -> Element) {
    self._base = _base
    self._transform = _transform
  }
}

因为实现了 IteratorProtocol，所以 LazyMapIterator 必须要实现 next() 方法，而关键点也是在 next() 方法中：

public mutating func next() -> Element? {
  return _base.next().map(_transform)
}

这里我们需要注意到的是 next() 方法中，调用的是 _base.next().map(_transform)，而不是 _base.next()，可以看出，map 的映射是读操作（访问序列元素时再做转换），这里面的 map 就是前面已经介绍过的 map() 方法，不重复说了。

我们再看 LazyMapCollection 源码（只保留最关键的 subscript() 方法）：

public struct LazyMapCollection<
  Base : Collection, Element
> : LazyCollectionProtocol, Collection {

  // ...
  
  @_inlineable
  public subscript(position: Base.Index) -> Element {
    return _transform(_base[position])
  }]
  
  // ...
  
}

这里可以看出，LazyMapCollection 也是提供一种读操作的机制，在使用的时候才真正地做转换，不使用则不做任何处理。

源码这里就不完整读下去了，建议读者自行去阅读下。接着我们来实践一下刚才讲解的源码，第一事当然是先看下 LazyMapSequence 文档：

F632EEC5-F40A-42BC-9C21-6F0E5F78167A.png

文档上写得很明显了，也不需要多解释，我们直接代码搞定：

var array: [Int] = [1, 2, 3]
print(array.lazy)     // LazyRandomAccessCollection<Array<String>>(_base: [1, 2, 3])
let arr = array.lazy.map({ $0 * 2 })
print(arr)  // LazyMapRandomAccessCollection<Array<Int>, Int>(_base: [1, 2, 3], _transform: (Function))
print(arr[0]) // 2

小结

在这篇文章里埋下了不少的坑，makeIterator()、next() 这些都是属于 Sequence 里面的内容了，这里都只是一提而过，所以下一次要分享的是 Sequence。