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。

推荐阅读
Swift 解读 - 前言
Swift 解读 - Collection 大家族(上篇)

推荐阅读更多精彩内容

  • 一、数组中的 map 和 flatMap 数组中的 map 对数组元素进行某种规则的转换,例如: 二、 flatM...
    kmplayer阅读 388评论 0 2
  • map 和 flatMap 是 Swift 中两个常用的函数,它们体现了 Swift 中很多的特性。对于简单的使用...
    SwiftCafe阅读 750评论 7 16
  • 粗鄙的简介 以下转自CSDN的一篇博博客http://blog.csdn.net/fish_yan_/articl...
    小白猿阅读 54评论 0 0
  • 6.1 唐巧烧脑体操3 总结 体操三 总结 mapmap 方法接受一个闭包作为参数, 然后它会遍历整个数组,并对数...
    婉卿容若阅读 2,284评论 0 4
  • 写在最前面声明:此篇博文非本人原创,是我在学习RxSwift框架时在网上找的,网上其他版本的阅读效果我不想吐槽,于...
    Qsyx阅读 28评论 0 0