SE-0008 | Swift Evolution↗

Add a Lazy flatMap for Sequences of Optionals

Proposal: SE-0008↗
Authors: Oisin Kidney
Review Manager: Doug Gregor
Status: Implemented (Swift 3.0)

01 何が問題だったのか

初期のSwiftの標準ライブラリには、flatMap が2種類用意されていました。ひとつはシーケンスのシーケンスを変換後に平坦化するもので、もうひとつは Optional のシーケンスから nil を取り除いて値だけを取り出すものです。

// ネストしたシーケンスを平坦化する flatMap
[1, 2, 3]
  .flatMap { n in n..<5 }
// [1, 2, 3, 4, 2, 3, 4, 3, 4]

// Optional のシーケンスを平坦化する flatMap
(1...10)
  .flatMap { n in n % 2 == 0 ? n / 2 : nil }
// [1, 2, 3, 4, 5]

問題は、遅延評価版（lazy を挟んだ版）がこのうち片方にしか用意されていなかったことです。ネストしたシーケンス向けの flatMap には遅延版がある一方、Optional を返すクロージャを受け取る版には遅延版がなく、lazy を挟んでも即時に評価された Array が返ってきてしまう状態でした。

// こちらは遅延版がある
[1, 2, 3]
  .lazy
  .flatMap { n in n..<5 }
// LazyCollection<FlattenBidirectionalCollection<...>>

// こちらは lazy を挟んでも即時評価されてしまう
(1...10)
  .lazy
  .flatMap { n in n % 2 == 0 ? n / 2 : nil }
// [1, 2, 3, 4, 5]

遅延版が無いことのデメリット

命令的にネストした for ループを、メソッドチェーンに書き換えたい場面はよくあります。このとき即時評価の flatMap を使うと、途中結果として不要な配列が毎段階で確保されてしまいます。ネストしたシーケンス向けの flatMap には遅延版が用意されているため中間配列の確保を避けられますが、Optional を返す変換を含むチェーンでは同じ最適化ができず、遅延評価の恩恵が途切れてしまっていました。

Optional を返す flatMap は、「変換と同時にフィルタリングを行う」という、チェーンの中でもとくに出番の多いパターンです。そこに遅延版が無いのは、標準ライブラリのAPIとしての一貫性の面でも穴でした。

02 どのように解決されるのか

LazySequenceType と LazyCollectionType（いずれも当時の名称。現在の LazySequenceProtocol / LazyCollectionProtocol に相当）に、Optional を返すクロージャを受け取る遅延版の flatMap を追加します。これにより、lazy を挟んだチェーンの中でも、中間配列を確保せずに「変換 + nil の除去」を行えるようになります。

使い方

使い方は既存の flatMap と同じで、lazy を挟んだシーケンス／コレクションに対して、Optional を返すクロージャを渡します。結果は遅延シーケンス／遅延コレクションとして返り、実際の要素は取り出されるタイミングで計算されます。

let results = (1...1_000_000)
  .lazy
  .flatMap { n -> Int? in
    n % 2 == 0 ? n / 2 : nil
  }

// この時点ではまだ計算されない
// 先頭 5 要素を取り出したタイミングで、必要な分だけ評価される
for x in results.prefix(5) {
  print(x)
}
// 1, 2, 3, 4, 5

lazy を付けない場合はこれまでどおり即時評価の flatMap が選ばれ、Array が返ります。使い分けは呼び出し側の lazy の有無だけで決まります。

実装の考え方

新しい型を導入するのではなく、標準ライブラリに既にある遅延版の map と filter の組み合わせで実現しています。概念的には次のようなチェーンと等価です。

extension LazySequenceType {
  public func flatMap<T>(transform: Elements.Generator.Element -> T?)
    -> LazyMapSequence<LazyFilterSequence<LazyMapSequence<Elements, T?>>, T> {
      return self
        .map(transform)
        .filter { opt in opt != nil }
        .map { notNil in notNil! }
  }
}

LazyCollectionType 向けにもほぼ同じ形で提供されます。一方、双方向にたどれるコレクション向けの遅延版（BidirectionalCollection に相当するもの）は、当時の標準ライブラリに FilterBidirectionalCollection が無いため提供されません。双方向版が欲しい場合は Array に一度取り出すなど、即時評価を挟む必要があります。

補足: `compactMap` への改称

なお、Optional を返すクロージャを受け取る flatMap は、のちに別提案（SE-0187, Swift 4.1）で compactMap に改称されました。現在のSwiftでは、本提案で追加された「遅延版」も含めて compactMap という名前で使います。

// 現在の書き方
let results = (1...1_000_000)
  .lazy
  .compactMap { n -> Int? in
    n % 2 == 0 ? n / 2 : nil
  }