map 、 flatMap

map 和 flatMap 是 Swift 中常用的高階函數(shù)，下面將從源碼的角度來解析使用，需要注意的是 map 和 flatMap 都是 Sequence 協(xié)議的擴展方法，并不是 Map.swift 里面的 map 方法，這兩個方法的源碼解讀全為了下文作鋪墊。

map 方法

map 接受一個閉包作為參數(shù)，作用于數(shù)組中的每個元素，然后將這些變換后的元素組成一個新的數(shù)組并返回：

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 協(xié)議的一個擴展方法，源碼也很簡單，直接對每個元素 transform，然后將 transform 的元素組成一個新的數(shù)組。

flatMap

關(guān)于 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 方法非常相似的，區(qū)別就是 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
  }
  

  這個方法的關(guān)鍵代碼是：

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

  如果數(shù)組中存在可選值或者 nil 時，系統(tǒng)就會調(diào)用這個重載，會把 optional 強制解包并把 nil 過濾掉：

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

通過閱讀 map 和 flatMap 的源碼，我們很易容就學(xué)會了這兩個高階函數(shù)的使用，鋪墊也說了，下面來到今天主要想說的內(nèi)容。

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 實現(xiàn)了 LazySequenceProtocol，而 LazySequenceProtocol 則是實現(xiàn)了 Sequence，同樣 LazyMapSequence 需要實現(xiàn) Sequence 協(xié)議，所以必需實現(xiàn) makeIterator 方法：

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

關(guān)鍵點是這里使用了新定義的 LazyMapIterator 迭代器，我們跳轉(zhuǎn)到 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
  }
}

因為實現(xiàn)了 IteratorProtocol，所以 LazyMapIterator 必須要實現(xiàn) next() 方法，而關(guān)鍵點也是在 next() 方法中：

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

這里我們需要注意到的是 next() 方法中，調(diào)用的是 _base.next().map(_transform)，而不是 _base.next()，可以看出，map 的映射是讀操作（訪問序列元素時再做轉(zhuǎn)換），這里面的 map 就是前面已經(jīng)介紹過的 map() 方法，不重復(fù)說了。

我們再看 LazyMapCollection 源碼（只保留最關(guān)鍵的 subscript() 方法）：

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

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

這里可以看出，LazyMapCollection 也是提供一種讀操作的機制，在使用的時候才真正地做轉(zhuǎn)換，不使用則不做任何處理。

源碼這里就不完整讀下去了，建議讀者自行去閱讀下。接著我們來實踐一下剛才講解的源碼，第一事當然是先看下 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

小結(jié)

在這篇文章里埋下了不少的坑，makeIterator()、next() 這些都是屬于 Sequence 里面的內(nèi)容了，這里都只是一提而過，所以下一次要分享的是 Sequence。

推薦閱讀
Swift 解讀 - 前言
Swift 解讀 - Collection 大家族（上篇）