前回(といっても大分前になってしまいましたが)、flatMap==flatten(map())であることから、flatMap()の準備としてflatten()を先に読みました。

Swiftのflatten()は平坦化を直ちに行うのではなく、要素が必要になった時にはじめて処理が実行されるよう、遅延評価の仕組みが取り入れられたのでした。この「元配列を別視点から見られるようにする」という仕組みによって、実行コストを抑えたまま平坦化を実現する面白いメソッドになっていることが分かりました。

そんなflattenの振る舞いを踏まえて、今回flatMap()に進みます。SwiftのflatMap()は、はたしてどのような実装になっているのでしょうか。

…とその前に、ここまでにmap,flattenを見てきたところとして残っている疑問がありますので、先に書き出しておきます。

flatten()の返す型はflatMap()のそれとは異なるが、いつ、どうやって変換しているか
flatten()のunderestimateCount()は常に0を返すけれど大丈夫？

どちらもflatMapを読んで解決できるといいのですが…。それでは読んでみましょう。

flatMap()を提供している型

flatMap()は、Optional, ImplicitlyUnwrappedOptional、SequenceTypeの3つの型で利用できます。ソースコードの定義箇所は、それぞれ以下のとおりです。

Optional : stdlib/public/core/Optional.swift
ImplicitlyUnwrappedOptional : stdlib/public/core/ImplicitlyUnwrappedOptional.swift
SequenceType : stdlib/public/core/SequenceAlgorithms.swift.gyb

さらには、遅延評価型のLazySequenceTypeとLazyCollectionTypeでも定義されているので、合計5箇所です。

LazySequenceType : stdlib/public/core/flatMap.swift
LazyCollectionType : stdlib/public/core/flatMap.swift

OptionalとImplicitlyUnwrappedOptionalのflatMap()

OptionalとImplicitlyUnwrappedOptionalの2つは、以前の記事で読んでいるので省略します。

swiftlife.hatenablog.jp

やっていることはとても単純で、アンラップしたオブジェクトに関数を適用して返すことをしていたのでした。

SequenceTypeのflatMap()

SequenceTypeのflatMap()は、SequenceAlgorithms.swift.gybの中で、SequenceTypeのextensionとして一般化された内容で記されています。2つあり、それぞれ短い内容です。

// 1つ目
@warn_unused_result
public func flatMap<S : SequenceType>(
  transform: (${GElement}) throws -> S
) rethrows -> [S.${GElement}] {
  var result: [S.${GElement}] = []
  for element in self {
    result.appendContentsOf(try transform(element))
  }
  return result
}

// 2つ目
@warn_unused_result
public func flatMap<T>(
  @noescape transform: (${GElement}) throws -> T?
) rethrows -> [T] {
  var result: [T] = []
  for element in self {
    if let newElement = try transform(element) {
      result.append(newElement)
    }
  }
  return result
}

ほとんど同じ内容ですが、1つ目がtransform引数の戻り型が「S」であることからnilを返さないこと、2つ目は逆に「T?」となっていて、Optionalで返すことができます。

ここでの${GElement}は、同gybファイルの先頭でGenerator.Elementの言い換えであると定義されているものです。

# We know we will eventually get a SequenceType.Element type.  Define
# a shorthand that we can use today.
GElement = "Generator.Element"

要素をtransform()に渡しつつ、その結果をresultにappend()あるいはappendContentsOf()で追加しています。言うなればflatMap()というよりconcatMap()と表現した方がしっくりくる感じです。

LazySequenceTypeとLazyCollectionTypeのflatMap()

flatMap.swiftという、そのものの名前のファイルでは、LazySequenceTypeとLazyCollectionTypeのflatMap()が定義されています。これらは関数宣言が複雑に見えますが、やっていることはどれも同じで次の1行が記述されているのみです。

return self.map(transform).flatten()

これが1番シンプルですね。flatMap()の由来である「flatten(map(x))」そのものになっています。

まとめ

これまで長々とflatMap()に至る実装を追いかけてきましたが、ようやく一区切りつけることができました。

flatten(map(x))が由来のflatMap()ですが、SwiftではLazySequenceTypeとLazyCollectionType以外はflatten()を使っておらず、それぞれの型で異なる実装になっていることがわかりました。実装を区別しているのは、その方がパフォーマンスが良いのかもしれませんし、あるいは疑問の一つだったflatten()の特性「underestimateCount()が0を返す」から、期待する結果をベースに実装を考えているのかもしれません。

細かな不明点・疑問点は残っていますが、それはおいおい調べていこうと思います。

おまけ

RxJS, RxSwiftなどReactive Programmingライブラリのドキュメントで、flatMapを表すとても分かりやすい図があったので引用します。

ReactiveX - FlatMap operator

配列（ストリーム）の要素に対して順に処理を実施し、処理では任意型の配列を返します。そして各結果をflattenして、最終的な別のストリームを返します。

swift life

（ようやく）SwiftのflatMap()を追いかける

flatMap()を提供している型

OptionalとImplicitlyUnwrappedOptionalのflatMap()

SequenceTypeのflatMap()

LazySequenceTypeとLazyCollectionTypeのflatMap()

まとめ

おまけ