Enumeratorを返すArray#flatten（Ruby）

@kts_h さんの
複数の Enumerable なオブジェクトで tally したい【追記あり】
の記事でコメントのやり取りをしていて、Array#flattenに Enumerator を返すメソッドがあるとおもしろいと思ったので、適当に実装してみました。

##コード

Ruby

class Array
  def flatten_to_enum
    Enumerator.new do |y|
      doit = ->(e) {
        if e.class == Array
          e.each {doit.(_1)}
        else
          y << e
        end
      }
      each {doit.(_1)}
    end
  end
end

##使い方

Ruby

ary = [0, 1, [2, [3], [4, 5, [6], 7, 8], 9, 10], 11]

ary.flatten
#=>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

ary.flatten_to_enum.to_a
#=>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
ary.flatten_to_enum.map {_1 ** 2}
#=>[0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121]

もし巨大な Array で複雑に入れ子になっているものがあったら、平坦化して利用するのにメモリを節約できます。ただ、この実装は再帰を使っているので、あまりにも深くネストしている配列は平坦化できないかも知れませんが、まあ実際上は大丈夫でしょう。

しかし、そんな需要があるのかというと、まずないだろうなあ（笑）。

##追記
@kts_h さんのコメントを受けてコードを修正しました。それから、上のリンク先の記事のような場合の例を挙げておきます。

Ruby

ary = Array.new(200000) {Array.new(20) {rand(5)}}
ary.flatten_to_enum.tally
#=>{2=>800022, 1=>799687, 4=>798559, 0=>800752, 3=>800980}

こうした巨大な Array がネストしている場合、Array#flattenだとまた巨大な配列を作ってしまってtallyすることになります。こうした場合にメモリ使用量の点で有用かなということです。Enumerator だし頻繁に Proc をネストして呼び出すので、速度は遅いです。

##再追記
Proc 呼び出しの回数が減るように修正しました。

Ruby

class Array
  def flatten_to_enum
    Enumerator.new do |y|
      store_array = ->(ary) {
        ary.each do |e|
          if e.class == Array
            store_array.(e)
          else
            y << e
          end
        end
      }
      store_array.(self)
    end
  end
end

こちらの方が速いと思います。