Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Memoization through Rubinius AST Transforms
# 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
"""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.