Memoization through Rubinius AST Transforms

memoization_ast_transform.rb

# This is a Rubinius AST Transform that implements a `memoization` keyword to
# be put before method definitions.
#
# Inspired by José Valim's talk at RuPy 2011.
#
# __KNOWN CAVEATS__: It doesn't reset the cache when changing arguments (it just
# memoizes the first call and returns that value forever more).
#
# Any ideas to make it work better? :D
#
class MemoizationTransform < Rubinius::AST::Send
  attr_accessor :meth

  transform :default, :memoize, "Memoization transform"

  def self.match?(line, receiver, name, arguments, privately)
    if name == :memoize
      m = arguments.body.first
      if m.is_a?(Rubinius::AST::Define)
        transform = new line, receiver, name, privately
        transform.meth = m
        transform
      end
    end
  end

  def bytecode(g)
    pos(g)

    meth.body.array = wrap(meth.body.array, meth.name)

    m = Rubinius::AST::Define.new(meth.line, meth.name, meth.body)
    m.arguments = meth.arguments
    m.bytecode(g)
  end

  private

  def wrap(instructions, name)
    o = Object.new
    o.instance_eval { @instructions = instructions; @name = name }

    def o.bytecode(g)
      # Check the cache
      g.push_ivar :"@__memoized_#{@name}"

      # If the method is already memoized, just go to the end and return the
      # memoized value.
      fin = g.new_label
      g.dup
      g.git fin

      # Otherwise, execute the method.
      @instructions.each do |b|
        b.bytecode(g)
      end

      # And memoize the value for further calls.
      g.set_ivar :"@__memoized_#{@name}"
      g.ret
      fin.set!
    end
    [o]
  end
end

Rubinius::AST::Transforms.register :default, :memoize, MemoizationTransform

eval """
  $times_called = 0

  memoize def foo(bar, baz)
    $times_called += 1
    bar + baz
  end

  p foo(1,3)
  p foo(1,3)
  p foo(1,3)
  p foo(9,12)
  p foo(1,3)
  p foo(1,3)

  p $times_called
"""