Fiber を使って Enumerable#lazy を再発明してみた

Ruby2.0 で追加されると噂の Enumerable#lazy は、Ruby で遅延リストが使えるようになる、個人的に待望の機能。map や select をメソッドチェーンしたとき、何度もループをぶん回すのが不満だったからなぁ。

Enumerable#lazy は以前ちらっとソースコードを見たんだけど、確か Enumerator を使って実装していた。Enumerator はソースコード見たこと無いけど、Fiber を使って実装したらしい。Fiber 使ったこと無いな、そういえば。

Fiber の勉強として、Enumerable#lazy を再発明してみた。あくまで勉強なので、Enumerator は使わず、Enumerator もどきを Fiber を使って自作している。

# 自作 Enumerator
class MyEnumerator
  include Enumerable

  # each を終了させるための例外
  class StopIteration < Exception; end

  # Fiber をラップして使いやすくする
  class Yielder
    def initialize(&block)
      @fiber = Fiber.new do
        block.call(self)
        raise StopIteration
      end
    end

    def <<(value)
      Fiber.yield(value)
    end

    def next
      @fiber.resume
    end
  end

  def initialize(collection=[], method=:each, &block)
    if block_given?
      @block = block
    else
      @block = Proc.new do |yielder|
        collection.send(method) do |args|
          yielder << args
        end
      end
    end
    @yielder = Yielder.new(&@block)
  rescue StopIteration
  end

  def next
    @yielder.next
  end

  # StopIteration が発生するまで値を取り出し続ける
  def each(&block)
    loop do
      block.call(self.next)
    end
  rescue StopIteration
  end
end

module Enumerable
  # 自作の Lazy クラス
  class MyLazy < MyEnumerator
    def initialize(collection, &block)
      super(collection, :each, &block)
    end

    # 遅延評価バージョンの map
    def map(&block)
      MyLazy.new(self) do |yielder|
        self.each do |n|
          yielder << block.call(n)
        end
      end
    end

    # 遅延評価バージョンの select
    def select(&block)
      MyLazy.new(self) do |yielder|
        self.each do |n|
          if block.call(n)
            yielder << n
          end
        end
      end
    end
  end

  # コレクションを MyLazy に変換
  def my_lazy
    MyLazy.new(self)
  end
end

# テスト
(1..10).my_lazy.map { |n|
  puts "map:#{n}"
  n * n
}.select { |n|
  puts "select:#{n}"
  n % 2 == 0
}.each { |n|
  puts "each:#{n}"
}

実行結果

map:1
select:1
map:2
select:4
each:4
map:3
select:9
map:4
select:16
each:16
map:5
select:25
map:6
select:36
each:36
map:7
select:49
map:8
select:64
each:64
map:9
select:81
map:10
select:100
each:100