素数列を enumerable ぽく作る

ruby

下から数えて1万番目の素数を求めよ、という何を求められているのか実はよく分からないコーディング課題に対してRubyで以下のように書いてみた。PrimeSeq[n]で、n番目の素数が得られる。

class PrimeSeq
  include Enumerable

  def initialize
    @primes = [2, 3]
    @i = 3
  end

  def [](idx)
    if @primes.size >= idx
      @primes[idx - 1]
    else
      take(idx - @primes.size).last
    end
  end

private

  def each
    loop do
      loop do
        @i += 2
        if @primes.none? {|p| (@i % p == 0)}
          @primes << @i
          break
        end
      end
      yield @i
    end
  end
end

if __FILE__ == $0
  prime = PrimeSeq.new
  (1..1000).each {|i| p prime[i]}
end

何となく、enumerable にすることで列っぽいものが扱えるだろうと思って、each を実装したけど、結局内部的に take のために使っているだけになった。ただ、列生成を再開すべきポインタを自分で意識的に保持しておかなくても「続きのm個を持って来い」と言えるのは良い、ように思うのだけど、いきなり self に対して take しているとか、しかもその take が暗黙的に private の each を呼んでるのだとか、そういう「Rubyだからね！」というやり方がいいのかどうか、よく分からないな。

そういえば、 MiniTestというのを使ってみた。しかし、RSpecとそんなに違うのかな、ていうか、また別のDSLなのか、ていうか、オレがやりたいのはassertなんだけど！ という感じに萎え萎え気分。RSpecより速いらしいし、Ruby1.9標準というのがいいのかな。

require './primeseq.rb'
require 'minitest/autorun'

describe PrimeSeq do
  before do
    @p = PrimeSeq.new
  end

  describe "from 1st to 5th" do
    it "should be" do
      @p[1].must_be :==, 2
      @p[2].must_be :==, 3
      @p[3].must_be :==, 5
      @p[4].must_be :==, 7
      @p[5].must_be :==, 11
    end
  end

  describe "1000th, 2000th" do
    it "should be" do
      @p[1_000].must_be :==, 7919
      @p[2_000].must_be :==, 17389
      @p[1_000].must_be :==, 7919
    end
  end

  describe "fetch from the cache" do
    it "should be" do
      @p[2_000].must_be :==, 17389
      @p[2_000].must_be :==, 17389
      @p[1_000].must_be :==, 7919
    end
  end
end