Swift标准库源码之旅 - LazySequence

背景

Lazy sequences can be used to avoid needless storage allocation and computation, because they use an underlying sequence for storage and compute their elements on demand.

在使用数组或某一序列的map/filter等方法的时候，容器会立刻进行遍历对所有元素执行transform和isIncluded闭包，时间复杂度是O(n)，并且返回的新数组也会占用对应内存空间。如果transform里的操作比较耗时，容器元素比较多，那么此种方式的时间消耗就不能忽略了。

举个极端的例子：

       let result1 = Array(repeating: 3, count: 99).map { (num) -> String in
            (0 ... 9999).reduce("0") { (current, num) in
                return current + num.description
            }
        }

上述代码的执行时间需要5.8s. 即使我们还根本没有用到其中的元素.

Lazy的方式就是将transform等需要遍历元素的操作进行延时，在需要元素的时候再进行transform.

思路

将transform闭包作为属性保存下来，重写迭代方法，在迭代之前元素的基础上进行·transform·返回。

实现 LazyMapSequence

所有的Sequence皆可进行lazy. 为Sequence类型添加扩展方法

extension Sequence {
 var lazy: LazySequence<Self> {
    return LazySequence(_base: self)
  }
}

LazySequence是一个中间类型，把原始Sequence作为属性保存起来备用.后面的map/filter都在这个类型基础上扩展

struct LazySequence<Base: Sequence>: Sequence {
    typealias Element = Base.Element
    
    var _base: Base
    init(_ base: Base) {
        self._base = base
    }
    
    func makeIterator() -> Base.Iterator {
        return _base.makeIterator()
    }
}

为LazySequence提供map方法，返回值是一个重写了迭代方法的新的Sequence

extension LazySequence {
    func map<ResultElement>(_ transform: @escaping (Element) -> ResultElement) -> LazyMapSequence<Base, ResultElement> {
        LazyMapSequence(self._base, transform: transform)
    }
}

他要把transform和原始序列保存下来.

struct LazyMapSequence<Base: Sequence, Element>: Sequence {
    
    typealias Element = Element
        
    var _base: Base
    var transform: (Base.Element) -> Element
    
    init(_ base: Base, transform: @escaping (Base.Element) -> Element) {
        self._base = base
        self.transform = transform
    }
    
    func makeIterator() -> Iterator {
        Iterator(_base.makeIterator(), transform: transform)
    }
}

他的迭代方法如下,取base的元素然后进行transform返回

extension LazyMapSequence {
    struct Iterator: MyIteratorProtocol {
        var _base: Base.Iterator
        var transform: (Base.Element) -> Element
        
        init(_ base: Base.Iterator, transform: @escaping (Base.Element) -> Element) {
            self._base = base
            self.transform = transform
        }
        
        mutating func next() -> Element? {
            return _base.next().map(transform)
        }
    }
}

以上仅解决了Sequence进行map延迟遍历的问题.
更常用的情况是一个数组map后使用下标取值.也需要支持.

也就是当LazyMapSequence的Base是个Collection的时候

public typealias LazyMapCollection<T: Collection,U> = LazyMapSequence<T,U>

然后就可以给LazyMapCollection类型扩展方法了

extension LazyMapCollection: Collection {
  ...
  ...
  subscript(position: Base.Index) -> Element {
    return _transform(_base[position])
  }
}

这样,数组类型的lazy后的map方法返回的LazyMapSequence就支持了下标访问.transform只在取某一值的时候调用.

上面只完成了map方法的lazy. 对于Filter方法想要实现lazy同样需要已基本相同的方式写一个LazyFilterSequence写其对应的迭代方法.

## 实现 LazyFilterSequence

extension LazySequence {
    func filter(_ isIncluded: @escaping (Element) -> Bool) -> LazyFilterSequence<Base> {
        LazyFilterSequence(_base, isIncluded: isIncluded)
    }
}

struct LazyFilterSequence<Base: MySequence>: MySequence {
    
    typealias Element = Base.Element
        
    var _base: Base
    var _predicate: (Base.Element) -> Bool
    
    init(_ base: Base, isIncluded: @escaping (Base.Element) -> Bool) {
        self._base = base
        self._predicate = isIncluded
    }
    
    func makeIterator() -> Iterator {
        Iterator(_base.makeIterator(), isIncluded: _predicate)
    }
}

extension LazyFilterSequence {
    struct Iterator: MyIteratorProtocol {
        var _base: Base.Iterator
        var _predicate: (Base.Element) -> Bool
        
        init(_ base: Base.Iterator, isIncluded: @escaping (Base.Element) -> Bool) {
            self._base = base
            self._predicate = isIncluded
        }
        
        mutating func next() -> Element? {
            while let value = _base.next() {
                if _predicate(value) {
                    return value
                }
            }
            return nil
        }
    }
}

与LazyMapSequence实现相似，不在赘述。

但是Filter是筛选元素得到一个新结果，数组的下标访问如何做到lazy呢? 不执行完Filter是不能拿到最后数组个数的吧。带着这个猜想发现了有意思的情况.

        let result = Array(repeating: 100, count: 99999).lazy.filter { $0 > 100 }
        print(result.count) // 0
        print(result[3]) // 100

这个筛选结果必然是空的.count为0是没问题的. 但是下标取值[3]确是返回了100

去看一下源码发现下标取值的方法是取的原数组的值, 并没有经过筛选

subscript(position: Index) -> Element {
    return _base[position]
}

所以下面这种操作方式也会得到你意想不到的结果.

        var array: [String?] = ["a", nil, nil, "c"]
        let result = array.lazy.filter { $0 != nil }
        for i in 0 ..< result.count {
            print(result[i])   // a, nil
        }

result本来是筛选过的不为nil的结果,期望值应该是"a"和"c". 结果却是"a"和"nil"

有一些方法没有新生成一个LazyxxSequence而是建立在上面两个Sequence之上的.

Lazy的CompactMap

func compactMap<ElementOfResult>(
    _ transform: @escaping (Elements.Element) -> ElementOfResult?
  ) -> LazyMapSequence<
    LazyFilterSequence<
      LazyMapSequence<Elements, ElementOfResult?>>,
    ElementOfResult
  > {
    return self.map(transform).filter { $0 != nil }.map { $0! }
  }

它是由LazyFilterSequence和LazyMapSequence组合实现的方法.因为用到了Filter. 下标取值仍然是取得原Collection的值.
但是自己实现时候发现LazyMapSequence类型其实并没有Filter方法, LazyFilterSequence方法也并没有map方法. 手动给具体类型加上吗? 但后面还有LazyDropWhileSequence类型,想要自由组合的话这种方式太硬了一点..
让他们共同遵循一个协议，给此协议添加扩展方法就可以让所有类型都具有map/filter

protocol LazySequenceProtocol: MySequence {
    
}

extension LazySequenceProtocol {
    func filter(_ isIncluded: @escaping (Element) -> Bool) -> LazyFilterSequence<Self> {
        LazyFilterSequence(self, isIncluded: isIncluded)
    }
}

extension LazySequenceProtocol {
    func map<U>(_ transform: @escaping (Element) -> U) -> LazyMapSequence<Self, U> {
        .init(self, transform: transform)
    }
}

接下来就有了上面CompactMap返回的多个尖括号的泛型嵌套类型。好处是可以自由组合。

为什么LazySequenceProtocol里没有定义方法，其实Swift源代码里是有的，但是笔者目前没有弄明白里面方法的意义。