【Ruby】竹内関数でのメモ化ふたたび

はじめに

竹内関数をメモ化で高速化する工程 の続編です。

今回は、既存のメモ化されていないメソッドないし関数のメモ化についてです。

前回同様、関数オブジェクトがある Python から見てみます。

Python の場合

勉強のため、自力でこんな関数を作っておくと、

def memoize(fn):
    cache = {}
    uniq = object()
    cache_get = cache.get
    def fun(*args):
        val = cache_get(args, uniq)
        if val is not uniq:
            return val
        val = fn(*args)
        cache[args] = val
        return val
    return fun

Python の場合、Ruby と違ってメソッドを実行するたびにメソッドオブジェクトを作り、それに引数を渡してメソッドオブジェクトを実行するようなので、cache_get に cache をレシーバにしたメソッドオブジェクトを入れておいて、それを使いまわすことで、効率的に計算できるようにしました。

その上で、

def tarai(x, y, z):
    if x <= y:
        return y
    val = tarai(
        tarai(x-1, y, z),
        tarai(y-1, z, x),
        tarai(z-1, x, y)
    )
    return val

tarai = memoize(tarai)

とするか、いっそのこと、アノテーションを使い、

@memoize
def tarai(x, y, z):
    if x <= y:
        return y
    val = tarai(
        tarai(x-1, y, z),
        tarai(y-1, z, x),
        tarai(z-1, x, y)
    )
    return val

することもできます。

オリジナルの tarai 関数は再帰関数です。ここで注意が必要なのは、再帰的に呼び出している tarai は今定義しつつあるtarai 関数オブジェクトではないということです。（局所関数の場合は事情が違ってくるかもしれません。多分、違います。）

この tarai は、今定義しつつある関数オブジェクトが実行される際に、tarai という名前が付いている関数オブジェクトなのです。要するに、新しい関数に tarai という名前を付けたら、元の tarai の中で実行される tarai は新しい 'tarai` だということなのです。

なので、tarai = memoize(tarai) すると、変数 tarai にはメモ化された関数が代入されています。この tarai 関数は、オリジナルの tarai などをキャプチャーしたクロージャです。ここで tarai(100, 50, 0) を実行すると、当然メモはまっさらなので、オリジナルの tarai 関数が呼び出されることになるのですが、その中で呼び出される tarai は、オリジナルの tarai （tarai という名前はすでに捨て去りました）ではなく、メモ化された現時点での tarai になるわけです。

では、Ruby に移ります。

Ruby の場合

まずは、lambda を使ったクロージャで挑戦（失敗）

lambda で tarai を作るとすると、こうでしょう。

def make_tarai
  tarai = lambda du |x, y, z|
    return y if x <= y
    tarai.call(
      tarai.call(x - 1, y, z),
      tarai.call(y - 1, z, x),
      tarai.call(z - 1, x, y)
    )
  end
end

tarai = make_tarai

そして、

def memoize(orig)
  cache = make_cache(3)
  lambda do |x, y, z|
    cache[x][y][z] ||= orig.call(x, y, z)
  end
end

tarai = memoize(tarai)

（make_cache は前回の記事を見てください。）
しかし、ちっとも高速化されません。理由はこうです。

lambda が伴っているブロック内では tarai.call により、Proc オブジェクトが再帰的に呼び出されています。しかし、この tarai は、Python の関数の場合のように、実行時に tarai という名前の付いたオブジェクトではないのです。つまり、オリジナルの tarai の中で呼び出される tarai は、tarai = make_tarai の tarai ではなく、tarai = lambda do ... の tarai なのです。

一度オリジナルの tarai に入り込んでしまうと、メモ化された tarai に戻ってくることができないのです。

残念。

メソッドで挑戦（まずは地道に）

まずは、メモ化されていないタライを作ります。

class Tarai
  def tarai(x, y, z)
    return y if x <= y
    tarai(
      tarai(x - 1, y, z),
      tarai(y - 1, z, x),
      tarai(z - 1, x, y)
    )
  end
end

tarai メソッドの呼び出しは、「tarai」という名前のメソッドを見つけて実行するということです。これなら、メモ化されたメソッドに 「tarai」という名前にすれば、オリジナルのたらいの中で呼び出される tarai はメモ化された tarai になるので上手くいきそうです。

次に、メモ化された tarai メソッドを作るのですが、その前に、オリジナルの tarai をどこかに退避しておく必要があります。Ruby の場合、メソッドに別名をつけることで退避できます。

こんな感じです。

class Tarai
  alias __tarai_orig tarai
  alias __init initialize

  def initialize
    __init
    @__cache = make_cache(3)
  end

  def tarai(x, y, z)
    @__cache[x][y][z] ||= __tarai_orig(x, y, z)
  end
end

した上で、

tarai = Tarai.new
tarai.tarai(100, 50, 0)

どうやら、上手くいったようです。

でも、メモ化するたびに決まり切ったコードを書くのはだるいので、前回同様、メタプログラミングで簡略化します。

メタプログラミングその1

def make_cache(arity)
  seed = "{}"
  (arity - 1).times do
    seed = "Hash.new { |h, k| h[k] = #{seed} }"
  end
  eval(seed)
end

def memoize(name)
  meth_name = name.to_s
  orig_name = "__#{meth_name}_orig"
  arity = instance_method(name.to_sym).arity
  args = arity.times.map { |i| "x#{i}" }.join(",")
  args2 = arity.times.map { |i| "[x#{i}]" }.join

  class_eval(<<~EOS)
    alias __init initialize
    alias #{orig_name} #{meth_name}
    def initialize(*args)
      __init(*args)
      @__cache = make_cache(#{arity})
    end
    def #{meth_name}(#{args})
      @__cache#{args2} ||= #{orig_name}(#{args})
    end
  EOS
end

class Tarai
  memoize \
  def tarai(x, y, z)
    return y if x <= y
    tarai(
      tarai(x - 1, y, z),
      tarai(y - 1, z, x),
      tarai(z - 1, x, y),
    )
  end
end

tarai = Tarai.new
tarai.tarai(100, 50, 0)

上手くいきました。アノテーションっぽくもできています。

メタプログラミングその2

オリジナルの tarai を別名で退避するのはかわいそうです。
そこで、発想を変えて、メモ化された tarai を Tarai オブジェクトのシングルトン・メソッドにして、オリジナルが欲しいときは super することにします。

module Memoization
  private \
  def make_cache(arity)
    seed = "{}"
    (arity - 1).times do
      seed = "Hash.new { |h, k| h[k] = #{seed} }"
    end
    eval(seed)
  end

  def memoize(name)
    arity = method(name.to_sym).arity
    args1 = arity.times.map { |i| "x#{i}" }.join(",")
    args2 = arity.times.map { |i| "[x#{i}]" }.join

    instance_eval(<<~EOS)
      @__cache = make_cache(#{arity})
      define_singleton_method(name) do |#{args1}|
        @__cache#{args2} ||= super(#{args1})
      end
    EOS
  end
end

class Tarai
  def tarai(x, y, z)
    return y if x <= y
    tarai(
      tarai(x - 1, y, z),
      tarai(y - 1, z, x),
      tarai(z - 1, x, y),
    )
  end
end

tarai = Tarai.new

tarai.extend Memoization
tarai.memoize :tarai

tarai.tarai(100, 50, 0)

def でメソッドを作らなかったためクロージャになったせいか、速度は落ちましたが、なんとか成功です。

まとめ

メタプログラミングをやりだすと楽しいですね。