Rubyで末尾再帰最適化をする。

元ネタはPythonで末尾再帰最適化をする。 - wasabizの日記。Pythonのデコレータを使って、末尾再帰で書かれた関数に対して末尾呼び出し最適化(TCO)を行う、というものです(どうやってTCOを実現しているかの詳細についての説明はここでは割愛します)。

さて、元エントリでは「Pythonがすごいからこんなことができるんだ」という感じで書かれていますが、タネさえ分かればいろんな言語でできそうだということが分かったので、他の言語でも試してみることにしました。

まずはじめに、試しにScheme版を書いてみたものの、そもそもSchemeはTCOを勝手にやってくれるのであまり意味のない例になってしまいました。
その後、Scheme版をだいたいそのままRubyに書き直したのが以下のRuby版です。

class Module
  def tco(name)
    continue = []
    first = true
    arguments = nil

    private_name = "private_" + name.to_s
    alias_method private_name, name
    private private_name

    proc = lambda do |*args|
      if first
        first = false
        while true
          result = send(private_name, *args)
          if result.equal? continue
            args = arguments
          else
            first = true
            return result
          end
        end
      else
        arguments = args
        continue
      end
    end
    define_method name, proc
  end
end

ここで定義した tco というクラスマクロを以下のように使います。

class Sum
  def sum1(n, acc=0)
    if n == 0
      acc
    else
      sum1(n-1, acc+n)
    end
  end

  def sum2(n, acc=0)
    if n == 0
      acc
    else
      sum2(n-1, acc+n)
    end
  end
  tco :sum2  # ←コレ
end

sum1メソッドとsum2メソッドは定義が全く同じですが、sum2メソッドにだけ tco をかけています。すると、以下のように、sum1メソッドでスタックオーバーフローになるようなケースでも、sum2メソッドではスタックオーバーフローになりません。TCOがうまく効いているようです。

>> o = Sum.new
=> #
>> o.sum1(100000)
SystemStackError: stack level too deep
        from ./tco.rb:36:in `sum1'
        from ./tco.rb:36:in `sum1'
        from (irb):3
>> o.sum2(100000)
=> 5000050000

実装としては、メソッドの呼び出しを置き換えるのに、Pythonのデコレータの代わりにアラウンドエイリアスを使ってます。アラウンドエイリアスの他にもメタプログラミングRubyで紹介されているテクニックをいくつか使わせてもらいました。メタプログラミングRuby素晴らしい！

メタプログラミングRuby