SwiftのLazySequence/LazyCollectionを使ってパフォーマンスを追求しよう

以下のコードでは2通りの方法でIntの配列を作っています。
どちらの方がパフォーマンスが良いでしょうか。

```Swift
let sequence = stride(from: 0, to: 1000000, by: 1)

// パターンA
let arrayA = sequence.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })

// パターンB
let lazySeq = sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })
let arrayB = Array(lazySeq)

arrayA == arrayB // => true
```

答えは、*パターンBの方がパフォーマンスが良い*です。
この記事ではその理由を解説します。

### お願い
注意して書いたつもりですが、間違いや勘違いなどあるかもしれません。その場合はご指摘お願いいたします:bow:

## 環境
* Swift 4.1
* macOS 10.13.4

## `map`や`filter`にはちょっと非効率なところがある
まずはパターンAの処理の詳細を考えてみます。
パターンAのほうが遅い理由は**無駄な配列を生成している**からです。

もう一度コードをみてみましょう。

```Swift:パターンAをもう一度
let sequence = stride(from: 0, to: 1000000, by: 1)
let arrayA = sequence.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })
```

処理の最初の部分は
`sequence.map({ $0 * ($0 + 2) })`
になっています。
[ドキュメント](https://developer.apple.com/documentation/swift/strideto/1689074-map)によると、`sequence.map({ $0 * ($0 + 2) })`は配列を返します。これが1つ目の配列です。

配列に対して`filter({ $0 % 2 == 0 })`を呼ぶともう一つ配列を作ります。
最後の`map({ $0 + 1 })`が3つ目の配列をつくります。

結局、配列を3つ作っています。
最終的に必要な配列は`arrayA`に代入するための1つだけで、残りの2つはすぐに解放されてしまいます。これは無駄です。
今回の例ではひとつひとつの配列の要素数が多いのためなおさらパフォーマンスに悪影響を与えます。

-----

対して、パターンBは配列を**1つ**しか作りません。
その理由をこれから解説します。

今度はパターンBのコードをもう一度確認しましょう。

```Swift:パターンBをもう一度
let sequence = stride(from: 0, to: 1000000, by: 1)
let lazySeq = sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })
let arrayB = Array(lazySeq)
```

途中にある`sequence.lazy`は
`LazySequence<StrideTo<Int>>`という型の値を返します。
こいつがキモになります。

## `LazySequence` / `LazyCollection`とは
記事の冒頭のコードはちょっとだけ忘れて、別の例で`LazySequence` / `LazyCollection`について解説します。
`Sequence`や`Collection`の`lazy`の典型的な使い方は重い処理を遅延させる処理です。

```swift
/// 画像を描画する重い処理
func heavyImageRender(_ index: Int) -> UIImage {
    // 重い処理
    return heavyImage
}

// 画像は描画されない
let lazyImageSequence = stride(from: 0, to: 100, by: 1).lazy.map(heavyImageRender)

// 1枚目のみ描画される
let firstImage = lazyImageSequence.first
```

コメントにもありますが、`lazyImageCollection.first`で1枚目の画像が本当に必要になったときになって初めて画像を描画します。
対して、`stride(from: 0, to: 100, by: 1).lazy.map(heavyImageRender)`から`lazy`を取り除くと100枚の画像を全て描画してしまいます。
このように、`Collection`や`Sequence`の要素が本当に必要になるまで評価しないメカニズムを提供してくれるのが`lazy`です( 「本当に必要になるまで評価しない」という処理は遅延評価ともよばれます )。

たとえば、Tinderのようなカードスワイプ系アプリで全ての画像をあらかじめロードしておくとパフォーマンスが悪くなってしまう場面などに効果を発揮します。

-----
さきほどのコードにある`lazyImageSequence`の型は`LazyMapSequence<StrideTo<Int>, UIImage>`です。
なんだかややこしい型です。最初の`LazyMapSequence`が`Sequence`に対して`map`の処理を遅延評価してくれるという意味になります。
後半の`<StrideTo<Int>, UIImage>`という部分は`Stride<Int>`がおおもとのSequenceで、最終的な評価の結果が`UIImage`であることを意味します。

他にもバリエーションがあります。

* `LazyFilterSequence`
  * `Sequence`に対して`filter`の処理を遅延評価してくれる
* `LazyMapCollection`
  * `Collection`に対して`map`の処理を遅延評価してくれる

いろいろなバリエーションがありますが、基本的には、`map`や`filter`といった処理を遅延させてくれます。
これらをSwiftでは[LazySequence](https://developer.apple.com/documentation/swift/lazysequence)や[LazyCollection](https://developer.apple.com/documentation/swift/lazycollection)と言うようです。

## `LazySequence` / `LazyCollection`の実装
もう一度この記事の冒頭のコードを確認します。

```swift:パターンB
let sequence = stride(from: 0, to: 1000000, by: 1)
let lazySeq = sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })
let arrayB = Array(lazySeq)
```

`sequence.lazy`の型は`LazySequence<StrideTo<Int>>`でした。
`LazySequence`の[実装](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/LazySequence.swift#L161-L184)を見ましょう。

重要な部分を抜粋します。

```swift:LazyCollection
public struct LazySequence<Base : Sequence>
: LazySequenceProtocol, _SequenceWrapper {
    internal init(_base: Base) {
        self._base = _base
    }

    public var _base: Base
}
```

`_base`というプロパティがあります。この`_base`には大元の`Sequence`が代入されます。この例では`sequence`つまり`stride(from: 0, to: 1000000, by: 1)`が代入されます。
____
次に、`sequence.lazy.map({ $0 * ($0 + 2) })`の型は`LazyMapSequence<StrideTo<Int>, Int>`です。
[`LazyMapSequence`の実装](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/Map.swift#L48-L88)をみてみましょう。

重要な部分を抜粋したのがこの以下のコードです。

```swift:LazyMapCollection
public struct LazyMapSequence<Base : Sequence, Element>
: LazySequenceProtocol {
  internal init(_base: Base, transform: @escaping (Base.Element) -> Element) {
    self._base = _base
    self._transform = transform
  }

  internal var _base: Base

  internal let _transform: (Base.Element) -> Element
}
```

プロパティに`_base`があるのは同じです。増えているのは`_transform`というプロパティです。
`_transform`プロパティはクロージャ式や関数への参照になっています。
したがって、

| プロパティ       |       代入される値 |
|:-----------------|:------------------|
| _base             |    `sequence.lazy`           |
| _transform           |            `{ $0 * ($0 + 2) }` |

になります。

-----
次に行きましょう。
`sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 })`は、`LazyFilterSequence<LazyMapSequence<StrideTo<Int>, Int>>`という型です。
[`LazyFilterSequence`の実装](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/Filter.swift#L63-L113)を見てみましょう。

重要な部分を抜粋するとこのようになっています。

```swift:LazyFilterSequence
public struct LazyFilterSequence<Base : Sequence>
: LazySequenceProtocol {
    public init(
        _base base: Base,
        _ isIncluded: @escaping (Base.Element) -> Bool
    ) {
        self.base = base
        self._include = isIncluded
    }

    internal var _base: Base
    internal let _predicate: (Base.Element) -> Bool
}
```

`LazyMapCollection`とよく似ています。表にするとこうなります。

| プロパティ       |       代入される値 |
|:-----------------|:------------------|
| _base             |    `sequence.lazy.map({ $0 * ($0 + 2) })`           |
| _predicate           |            `{ $0 % 2 == 0 })` |

----
最後は、
`let lazySeq = sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })`
です。
`lazySeq`は`LazyMapSequence`なので、以下のようになっています。

| プロパティ       |       代入される値 |
|:-----------------|:------------------|
| _base             |    `sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 })`     |
| _transform           |  `{ $0 + 1 }` |

-----
**ここまでの処理は非常に軽い**ことに注意してください。
やっていることはプロパティにクロージャ式を代入しているだけです。

### 最後に配列を生成する
配列を生成するコードは以下です。単にイニシャライザを呼んでいるだけです。

```swift
let arrayB = Array(lazySeq)
```

`Array`のイニシャライザは[ここの部分](https://github.com/apple/swift/blob/master/stdlib/public/core/ContiguousArrayBuffer.swift#L583-L607)がキモのようです。
イニシャライザの引数(ここでは`lazySeq`)に`makeIterator()`を読んでやってその結果をバッファにどんどんappendしていきます。

```swift
var iterator = source.makeIterator()

for _ in 0..<initialCapacity {
    builder.addWithExistingCapacity(iterator.next()!)
}

while let element = iterator.next() {
    builder.add(element)
}
```

`LazyMapSequence`は`makeIterator()`に[独自の実装](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/Map.swift#L62-L64)を与えています。

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

`_base.makeIterator` という部分の`_base`の型は`LazyFilterSequence`でした。
`LazyFilterSequence`にも`makeIterator()`の[独自の実装](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/Filter.swift#L76-L79)があります。

```swift:LazyFilterSequence
public func makeIterator() -> LazyFilterIterator<Base.Iterator> {
    return LazyFilterIterator(
      _base: base.makeIterator(), _include)
}
```

それぞれの`makeIterator()`は専用の`Iterator`として[`LazyMayIterator`](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/Map.swift#L13-L46), [`LazyFilterIterator`](https://github.com/apple/swift/blob/swift-4.1-branch/stdlib/public/core/Filter.swift#L13-L61)というのを返しています。
抜粋すると、

```swift:LazyMapIterator
public struct LazyMapIterator<
    Base : IteratorProtocol, Element> : IteratorProtocol, Sequence {
    public mutating func next() -> Element? {
        return _base.next().map(_transform)
    }
    
    public var base: Base { return _base }
    
    internal var _base: Base
    internal let _transform: (Base.Element) -> Element
    
    internal init(_base: Base, _transform: @escaping (Base.Element) -> Element) {
        self._base = _base
        self._transform = _transform
    }
}
```

```swift:LazyFilterIterator
public struct LazyFilterIterator<
  Base : IteratorProtocol
> : IteratorProtocol, Sequence {
    public mutating func next() -> Base.Element? {
        while let n = _base.next() {
            if _predicate(n) {
                return n
            }
        }
        return nil
    }

    internal init(
        _base: Base,
        _ isIncluded: @escaping (Base.Element) -> Bool
        ) {
        self._base = _base
        self._predicate = isIncluded
    }

    public var base: Base { return _base }

    internal var _base: Base

    internal let _predicate: (Base.Element) -> Bool
}
```

-----

これで登場する全ての型の紹介ができました。

`Array`のイニシャライザは`makeIterator()`を
`sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })`
に対して呼びます。すると、

* `LazyMapSequence`の`makeIterator()`の実装により`_base`( 型は`LazyFilterSequence` )に対してさらに`makeIterator()`を呼びます。
* `LazyFilterSequence`の`makeIterator()`の実装によりまたさらに`_base`( 型は`LazyMapSequence` )に対し`makeIterator()`を呼びます。
* `LazyMapSequence`の`makeIterator()`の実装によりまたさらに`_base`つまり、`sequence.lazy`に対して`makeIterator()`を呼びます。
* `sequence.lazy`は`_base`に`makeIterator()`を呼びます。[^1]
* `sequence`つまり`stride(from: 0, to: 1000000, by: 1)`は普通の`Sequence`なので、`makeIterator()`を呼ぶと普通の`Iterator`を返します。

Arrayのイニシャライザはこうしてできた`LazyMapSequence`に`next()`を次々と呼んでいきます。

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

`_base`つまり`LazyFilterIterator`に`next()`を呼び、その結果にクロージャ式を与えています。

```swift:LazyFilterIterator
public mutating func next() -> Base.Element? {
    while let n = _base.next() {
        if _predicate(n) {
            return n
        }
    }
    return nil
}
```

`LazyFilterIterator`の`next()`は`_base`の`next()`に対してクロージャを当てはめてtrueかfalseかを判断しているだけです。`LazyFilterIterator`の`_base`は同じく`LazyMapSequence`です。
こうしてさかのぼっていくと、
`sequence`つまり`stride(from: 0, to: 1000000, by: 1)`の`Iterator`にたどりつきます。こいつは`next()`を呼ぶたびに`0, 1, 2, 3, 4, ..., 999999, nil`を返します。

以上から、`Array`のイニシャライザはおおもとの`sequence`つまり`stride(from: 0, to: 1000000, by: 1)`をイテレートしながらつぎつぎにクロージャ式を呼んでいき、最後に一つの配列を作ります。
やっと、配列が1つしか作られない理由がわかりました！

## 最後にベンチマーク

以下のコードでベンチマークしました。マシンは`MacBook Pro (Retina, 13-inch, Mid 2014)`[^2]です。

```bash
$ swiftc --version
Apple Swift version 4.1 (swiftlang-902.0.48 clang-902.0.37.1)
Target: x86_64-apple-darwin17.5.0

$ cat patternA.swift
let sequence = stride(from: 0, to: 1000000, by: 1)
let arrayA = sequence.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })
print(arrayA.count, arrayA[100])

$ cat patternB.swift 
let sequence = stride(from: 0, to: 1000000, by: 1)
let lazySeq = sequence.lazy.map({ $0 * ($0 + 2) }).filter({ $0 % 2 == 0 }).map({ $0 + 1 })
let arrayB = Array(lazySeq)
print(arrayB.count, arrayB[100])

$ swiftc -O patternA.swift -o patternA
$ swiftc -O patternB.swift -o patternB

$ time ./patternA
500000 40401

real    0m0.101s
user    0m0.027s
sys     0m0.056s

$ time ./patternB
500000 40401

real    0m0.050s
user    0m0.033s
sys     0m0.013s
```

大幅に短くなっています :tada:

-----
[^1]: https://github.com/apple/swift/blob/a7ff0da33488b9050cf83df95f46e5b9aa2348d5/stdlib/public/core/SequenceWrapper.swift#L45-L47
[^2]: そろそろ新しいのを買いたいです :moneybag: :computer: