Ruby
Implementing a Macro in Ruby for Memoization
One of my favorite ways to learn is digging into existing libraries and trying to reimplement the functionality.
We were recently implementing our own memoization solution at work which led me to checking out some of the existing solutions out there. Memoist2[1] is a small implementation and I thought would be fun to walk through a similar implementation to see how it works.
For the uninitiated, "memoization" is:
an optimization technique used primarily to speed up computer programs by storing the results of expensive function calls and returning the cached result when the same inputs occur again. [2]
For example:
def foo
@foo ||= expensive_method
end
or
def foo
@foo ||= begin
expensive_stuff
end
end
We'll create a MemoRedux module that is included to provide
macros named memoize
and memoize_class_method
, which
will be used like this:
class Foo
include MemoRedux
def bar
end
memoize :bar
def self.baz
end
memoize_class_method :baz
end
(A ‘macro’ in this context is a class-level method that generates code.)
Here's a pretty similar looking API to the aforementioned gems. The first step will be to add the method definitions. We'll make it so multiple methods can be memoized at once.
module MemoRedux
def self.included(klass)
klass.extend(Macros)
end
module Macros
def memoize(*methods)
end
def memoize_class_method(*methods)
end
end
end
The next step will be to utilize Module#prepend
. The Module#prepend
method was added to Ruby 2.0 [3], and will add the module before it hits the current class, as opposed to
the typical Module#include
, which puts the module above the class in the call chain.
We'll get or set a module called MemoMod + the object id of self. Here self will be the class where
MemoRedux is included. This will help us avoid collisions with the random chance that a dynamically
created module has the same name as an existing constant somewhere else.
At the end of the memoize
method we
prepend this module.
pry(main)> module A; end
pry(main)> module C; end
pry(main)> module B; prepend A; include C; end
pry(main)> B.ancestors
=> [A, B, C]
module MemoRedux
def self.included(klass)
klass.extend(Macros)
end
module Macros
def memoize(*methods)
const_defined?(:"MemoMod#{self.object_id}") ?
const_get(:"MemoMod#{self.object_id}") : const_set(:"MemoMod#{self.object_id}", Module.new)
mod = const_get(:"MemoMod#{self.object_id}")
prepend mod
end
def memoize_class_method(*methods)
end
end
end
We'll loop through the methods to memoize, and call define_method for each one to create a new method to read the cached value.
module MemoRedux
def self.included(klass)
klass.extend(Macros)
end
module Macros
def memoize(*methods)
const_defined?(:"MemoMod#{self.object_id}") ?
const_get(:"MemoMod#{self.object_id}") : const_set(:"MemoMod#{self.object_id}", Module.new)
mod = const_get(:"MemoMod#{self.object_id}")
mod.class_eval do
methods.each do |method|
define_method(method) do
end
end
end
prepend mod
end
def memoize_class_method(*methods)
end
end
end
Inside define_method
, we'll create a hash to store the
computed values for the memoized methods, or call super()
to get the original method.
module MemoRedux
def self.included(klass)
klass.extend(Macros)
end
module Macros
def memoize(*methods)
const_defined?(:"MemoMod#{self.object_id}") ?
const_get(:"MemoMod#{self.object_id}") : const_set(:"MemoMod#{self.object_id}", Module.new)
mod = const_get(:"MemoMod#{self.object_id}")
mod.class_eval do
methods.each do |method|
define_method(method) do
@_memo_methods ||= {}
if @_memo_methods.include?(method)
@_memo_methods[method]
else
@_memo_methods[method] = super()
end
end
end
end
prepend mod
end
def memoize_class_method(*methods)
end
end
end
We'll fill out the memoize_class_method
next. We open
up the eigenclass by sending class_eval
to singleton_class
. This creates a
closure so we can use the method arguments *methods in the eigenclass.
MemoRedux is included here and we use the memoize
method again.
module MemoRedux
def self.included(klass)
klass.extend(Macros)
end
module Macros
def memoize(*methods)
const_defined?(:"MemoMod#{self.object_id}") ?
const_get(:"MemoMod#{self.object_id}") : const_set(:"MemoMod#{self.object_id}", Module.new)
mod = const_get(:"MemoMod#{self.object_id}")
mod.class_eval do
methods.each do |method|
define_method(method) do
@_memo_methods ||= {}
if @_memo_methods.include?(method)
@_memo_methods[method]
else
@_memo_methods[method] = super()
end
end
end
end
prepend mod
end
def memoize_class_method(*methods)
singleton_class.class_eval do
include MemoRedux
memoize(*methods)
end
end
end
end
It would be nice to be able to get a fresh value if we want,
so we'll add the ability to pass true
to skip the cached method.
instance.foo(true)
To do this we'll add an optional arg, skip_cache
, to the
dynamic method definition and add a check for it in the if
statement.
module Memoist
def self.included(klass)
klass.extend(Macros)
end
module Macros
def memoize(*methods)
const_defined?(:"MemoMod#{self.object_id}") ?
const_get(:"MemoMod#{self.object_id}") : const_set(:"MemoMod#{self.object_id}", Module.new)
mod = const_get(:"MemoMod#{self.object_id}")
mod.class_eval do
methods.each do |method|
define_method(method) do |skip_cache = false|
@_memo_methods ||= {}
if @_memo_methods.include?(method) && !skip_cache
@_memo_methods[method]
else
@_memo_methods[method] = super()
end
end
end
end
prepend mod
end
def memoize_class_method(*methods)
singleton_class.class_eval do
include Memoist
memoize(*methods)
end
end
end
end
And with this, we've arrived at the final solution. You can memoize instance methods, class methods, cache methods that return nil, and skip the cached method call.