Memoize（メモ化）モジュールを作ってみた

素数列とかフィボナッチ数列、ルーカス数列をクラス化して遊んでたら、「メモ化部分ってモジュール化できね？」と思いついた。
色々調べてみると、（Rubyで既に実装されたものも見つかったが、）PythonのMemoizeデコレータが便利そうなのでそれを意識してみた。
参考：Re: Pythonのデコレータが面白いのでRubyで実装してみた

ただそれだけじゃ面白くないので、メモ化宣言時に「初期値のリスト」を渡して計算を省略できるようにした。
すると、Haskellの

fibonacci 0 = 0
fibonacci 1 = 1
fibonacci n = fibonacci(n-2) + fibonacci(n-1)

のように、純粋にロジックのみの定義ができるようになった。これなら面白いかも。

使い方

クラスまたはモジュールにextendし、メモ化したいメソッドの直前でmemoize宣言*1をしてください。
その際、「引数の配列 => 値」のHashを渡すことで初期値を設定できます。引数が1つの場合は配列ではなく単独の値でも大丈夫です。
インスタンスメソッド、クラスメソッドどちらでも使用できますが、クラスメソッドに使用する際は下記のようにクラスの特異メソッド内でextendしてください。

フィボナッチ数列の例

module Fibonacci
  class << self
    extend Memoize
    
    memoize 0 => 0, 1 => 1
    def [] n
      self[n-2] + self[n-1]
    end
  end
end

p Fibonacci[0]  #=> 0
p Fibonacci[1]  #=> 1
p Fibonacci[2]  #=> 1
p Fibonacci[3]  #=> 2
p Fibonacci[10] #=> 55

ソース

module Memoize
  def self.extended mod
    mod.module_eval do
      class << self
        @memoized = {}
        
        def memoize default = {}
          @memoize_default = default.dup
          default.keys.each do |key|
            @memoize_default[[key]] = @memoize_default.delete(key) unless key.kind_of? Array
          end
        end
        
        private
        def method_added method
          return unless @memoize_default
          
          method_without = :"#{method}_without_memoize"
          memoized = (class << self; @memoized; end)
          memoized[method] = @memoize_default
          @memoize_default = nil
          
          alias_method method_without, method
          define_method method do |*args, &block|
            return __send__ method_without, *args, &block if block
            memoized[method][args] ||= __send__ method_without, *args
          end
        end
      end
      
      private
      def singleton_method_added method
        (class << self; self; end).__send__ :method_added, method
      end
    end
  end
end

既知の問題点

トップレベルへのextendができない。

class << self
  extend Memoize
  
  memoize 0 => 0, 1 => 1
  def fibonacci n
    fibonacci(n-2) + fibonacci(n-1)
  end
end
p fibonacci(10) #=> 55

どうしてもトップレベルで使いたいときはmainの特異メソッドに対してextendしてください。

module_functionとの相性が悪い。

引数無しのmodule_functionを使わず、対象のメソッドを定義した後に

module_function :fibonacchi, :fibonacchi_without_memoize

のようにしてください（ダサッ！！）。

補足

宣言みたいにしたかったからmethod_added系を使う黒魔術になってるけど、単純にメソッドを定義した後にメソッド名を指定してmemoizeを実行するだけだったらもっと綺麗になってるんだろうなぁ……

*1:実際にはただのメソッドですが