Swift標(biāo)準(zhǔn)庫源碼之旅 - 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.

在使用數(shù)組或某一序列的map/filter等方法的時(shí)候，容器會(huì)立刻進(jìn)行遍歷對(duì)所有元素執(zhí)行transform和isIncluded閉包竹挡，時(shí)間復(fù)雜度是O(n)镀娶，并且返回的新數(shù)組也會(huì)占用對(duì)應(yīng)內(nèi)存空間。如果transform里的操作比較耗時(shí)揪罕，容器元素比較多梯码，那么此種方式的時(shí)間消耗就不能忽略了。

舉個(gè)極端的例子：

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

上述代碼的執(zhí)行時(shí)間需要5.8s. 即使我們還根本沒有用到其中的元素.

Lazy的方式就是將transform等需要遍歷元素的操作進(jìn)行延時(shí)好啰，在需要元素的時(shí)候再進(jìn)行transform.

思路

將transform閉包作為屬性保存下來轩娶，重寫迭代方法，在迭代之前元素的基礎(chǔ)上進(jìn)行·transform·返回框往。

實(shí)現(xiàn) LazyMapSequence

所有的Sequence皆可進(jìn)行lazy. 為Sequence類型添加擴(kuò)展方法

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

LazySequence是一個(gè)中間類型鳄抒，把原始Sequence作為屬性保存起來備用.后面的map/filter都在這個(gè)類型基礎(chǔ)上擴(kuò)展

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方法，返回值是一個(gè)重寫了迭代方法的新的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的元素然后進(jìn)行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進(jìn)行map延遲遍歷的問題.
更常用的情況是一個(gè)數(shù)組map后使用下標(biāo)取值.也需要支持.

也就是當(dāng)LazyMapSequence的Base是個(gè)Collection的時(shí)候

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

然后就可以給LazyMapCollection類型擴(kuò)展方法了

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

這樣,數(shù)組類型的lazy后的map方法返回的LazyMapSequence就支持了下標(biāo)訪問.transform只在取某一值的時(shí)候調(diào)用.

上面只完成了map方法的lazy. 對(duì)于Filter方法想要實(shí)現(xiàn)lazy同樣需要已基本相同的方式寫一個(gè)LazyFilterSequence寫其對(duì)應(yīng)的迭代方法.

## 實(shí)現(xiàn) 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實(shí)現(xiàn)相似椰弊，不在贅述许溅。

但是Filter是篩選元素得到一個(gè)新結(jié)果，數(shù)組的下標(biāo)訪問如何做到lazy呢? 不執(zhí)行完Filter是不能拿到最后數(shù)組個(gè)數(shù)的吧秉版。帶著這個(gè)猜想發(fā)現(xiàn)了有意思的情況.

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

這個(gè)篩選結(jié)果必然是空的.count為0是沒問題的. 但是下標(biāo)取值[3]確是返回了100

去看一下源碼發(fā)現(xiàn)下標(biāo)取值的方法是取的原數(shù)組的值, 并沒有經(jīng)過篩選

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

所以下面這種操作方式也會(huì)得到你意想不到的結(jié)果.

        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的結(jié)果,期望值應(yīng)該是"a"和"c". 結(jié)果卻是"a"和"nil"

有一些方法沒有新生成一個(gè)LazyxxSequence而是建立在上面兩個(gè)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組合實(shí)現(xiàn)的方法.因?yàn)橛玫搅薋ilter. 下標(biāo)取值仍然是取得原Collection的值.
但是自己實(shí)現(xiàn)時(shí)候發(fā)現(xiàn)LazyMapSequence類型其實(shí)并沒有Filter方法, LazyFilterSequence方法也并沒有map方法. 手動(dòng)給具體類型加上嗎? 但后面還有LazyDropWhileSequence類型,想要自由組合的話這種方式太硬了一點(diǎn)..
讓他們共同遵循一個(gè)協(xié)議贤重，給此協(xié)議添加擴(kuò)展方法就可以讓所有類型都具有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返回的多個(gè)尖括號(hào)的泛型嵌套類型。好處是可以自由組合清焕。

為什么LazySequenceProtocol里沒有定義方法并蝗，其實(shí)Swift源代碼里是有的，但是筆者目前沒有弄明白里面方法的意義耐朴。