flatMap()の前にflatten()を調べる（予想外に深かった）

swift

Swiftを理解するための勉強中です。

前回、前々回でmap()を読んだので、次はflatMap()…といきたいところですが、flatMap()はmap+flattenなので、先にflattenの実装を読むことにします。

swiftlife.hatenablog.jp swiftlife.hatenablog.jp

flatten()メソッドは、Rubyのflatten、flatten!と同様、配列の配列を平坦化するためのメソッドです。Swiftだと、自身の要素を繋いだ(concatenate)ものと説明されています。flatten()が使用できるのは、 CollectionType,SequenceType,LazyCollectionType,LazySequenceTypeの4種類です。

lazy系は後日きちんと調べるので、今回はCollectionTypeのみ見てみましょう。すると、CollectionTypeにおけるflattenは、extensionで次のように実装されていました。

extension CollectionType where Generator.Element : CollectionType {
    /// A concatenation of the elements of `self`.
    @warn_unused_result
    public func flatten() -> FlattenCollection<Self>
}

extension CollectionType where Generator.Element : CollectionType, Index : BidirectionalIndexType, Generator.Element.Index : BidirectionalIndexType {
    /// A concatenation of the elements of `self`.
    @warn_unused_result
    public func flatten() -> FlattenBidirectionalCollection<Self>
}

おや、てっきりArrayを返すのかと思いきや、「FlattenCollection」「FlattenBidirectionalCollection」という初めて見る型を返しています。これらの型はいったい何でしょうか？ひとまずマニュアルを読んでみます。

struct FlattenCollection<Base : CollectionType where Base.Generator.Element : CollectionType> { ... }

struct FlattenBidirectionalCollection<Base : CollectionType where Base.Generator.Element : CollectionType, Base.Index : BidirectionalIndexType, Base.Generator.Element.Index : BidirectionalIndexType> { ... }

A flattened view of a base collection-of-collections.

The elements of this view are a concatenation of the elements of each collection in the base.

The flatten property is always lazy, but does not implicitly confer laziness on algorithms applied to its result. In other words, for ordinary collections c:

  • c.flatten() does not create new storage
  • c.flatten().map(f) maps eagerly and returns a new array
  • c.lazy.flatten().map(f) maps lazily and returns a LazyMapCollection

どちらの型にも同じ説明が書かれていて、「元コレクションを平坦化したときのビュー」(意訳)とあります。ふむふむ、これらは遅延評価されることが前提の型で、あくまで元のコレクションに対するビューポイント(視点)を提供するのみだと。その代わりに、実際に使われるまでは余計なメモリを消費しないと。どちらもCollectionTypeのサブクラスなのですね。

実装はファイルstdlib/public/core/Flatten.swift.gybにあるので、続いてこちらを見てみることにします。flattenはgybファイルを使い汎化された内容で書かれているので、直接のソースコードとしては出てきませんが、たぶんこれがCollectionTypeのflatten()メソッドと思われる記述がありました。次のソースです。

extension CollectionType where ${constraints % {'Base': ''}} {
  /// A concatenation of the elements of `self`.
  @warn_unused_result
  public func flatten() -> ${Collection}<Self> {
    return ${Collection}(self)
  }
}

自分自身(CollectionType)を元に${Collection}、つまりFlattenCollectionやFlattenBidirectionalCollectionをインスタンス化しています。なんとたったこれだけでした。${Collection}が指すイニシャライザでも、_baseプロパティにCollectionTypeオブジェクトを代入しているだけです。予想外の何もしてなさにびっくりしてしまいましたが、どうやらflatten()の秘密は、このFlattenCollection/FlattenBidirectionalCollection型が持つメソッドにあるようです。

FlattenCollection/FlattenBidirectionalCollectionが持つ、init()以外のメソッドは次の5種類です。

  • public func generate() -> FlattenGenerator<Base.Generator>
  • public var startIndex: FlattenCollectionIndex<Base> or FlattenBidirectionalCollectionIndex<Base>
  • public var endIndex: FlattenCollectionIndex<Base> or FlattenBidirectionalCollectionIndex<Base>
  • public subscript (position: FlattenCollectionIndex<Base> or FlattenBidirectionalCollectionIndex<Base>) -> Base.Generator.Element.Generator.Element
  • public func underestimateCount() -> Int

それぞれを読んでみたら（どれもstdlib/public/core/Flatten.swift.gybです）、ようやくflatten()の正体が見えてきました。

まずgenerate()の実装はreturn FlattenGenerator(_base.generate())とあり、FlattenGenerator型のオブジェクトを作成、返しています。このFlattenGeneratorはnext()メソッドのみ持っていて、ここで平坦化を実施しています。

startIndex()とendIndex()は、共にFlattenCollectionIndexまたはFlattenBidirectionalCollectionIndexを返し、そのsuccessor()/predecessor()メソッドによって平坦化した要素を返します。

そしてsubscriptも同じです。平坦化した状態でのpositionに相当する要素を返しています。なるほど、利用される可能性のあるメソッドすべてで、平坦化したときの要素を返すよう外堀を固めているのですね。もしflatten()を直ちに評価している場合、配列の要素をその場で全コピーすることになるので、コストは高くつきそうです。それにくらべてflatten()はこうして遅延評価を実現し、コストを最小限に抑えているんですね、おもしろいです。

ただ1点、underestimateCount()メソッドだけが次のような実装になっているのはまだ完全に理解できていません。

public func underestimateCount() -> Int {
  // To return any estimate of the number of elements, we have to start
  // evaluating the collections.  That is a bad default for `flatMap()`, so
  // just return zero.
  return 0
}

意訳すると「すべてを評価して数を返すというunderestimateCount()の振る舞いは、flatMap()にとっては都合が悪いので、常に0を返す」となりますが、逆にこの振る舞いで問題になったりしないのでしょうか？遅延評価のためだとはいえ、謎です…。

ともあれ、以上でflatten()の振る舞いをおおよそ把握することができました。はじめは単に平坦化した配列を返すだけと思っていたのですが、蓋を開けてみると別のオブジェクトを返しているだけで、遅延評価の外堀によって実行コストを抑える仕組みが整っているだけでした。とても面白いメソッドでした。

遅延評価、自分ももっと使えるようになりたいです。

参考サイト