Hashrocket.com / blog

Large macro

Implementing a Macro in Ruby for Memoization

posted on and written by in

Image 100x100 nick palaniuk

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.

References

  1. https://github.com/matthewrudy/memoist2
  2. https://en.wikipedia.org/wiki/Memoization
  3. http://gshutler.com/2013/04/ruby-2-module-prepend/

Posted in Ruby