"Achieving Fast Method Metaprogramming: Lessons from MemoWise" from RubyConf 2021

rubyconf_benchmarks.rb

# frozen_string_literal: true

require "json"
require "tempfile"
require "set"

require "benchmark/ips"

# Welcome to the benchmarks for "Achieving Fast Method Metaprogramming: Lessons
# from MemoWise" by Jacob Evelyn and Jemma Issroff, presented at RubyConf 2021.
#
# Note: some method definitions are the same between different versions. Rather
# than use inheritance, we use shared behavior to define them in both versions
# to avoid the chance that differences in the inheritance chain impact our
# results.

NO_ARGS_METHODS = %i[no_args no_args_falsey]
MEMOIZED_METHODS = NO_ARGS_METHODS + %i[
  one_positional_arg
  positional_args
  one_keyword_arg
  keyword_args
  positional_and_keyword_args
  positional_and_splat_args
  keyword_and_double_splat_args
  positional_splat_keyword_and_double_splat_args
]

def define_methods_for_testing_memo_wise(target)
  target.module_eval <<~END_OF_METHODS
    def no_args ; 100 ; end
    def no_args_falsey ; nil ; end
    def one_positional_arg(a) ; 100 if a.positive? ; end
    def positional_args(a, b) ; 100 if a.positive? ; end
    def one_keyword_arg(a:) ; 100 if a.positive? ; end
    def keyword_args(a:, b:) ; 100 if a.positive? ; end
    def positional_and_keyword_args(a, b:) ; 100 if a.positive? ; end
    def positional_and_splat_args(a, *args) ; 100 if a.positive? ; end
    def keyword_and_double_splat_args(a:, **kwargs) ; 100 if a.positive? ; end
    def positional_splat_keyword_and_double_splat_args(a, *args, b:, **kwargs) ; 100 if a.positive? ; end
  END_OF_METHODS

  MEMOIZED_METHODS.each do |method_name|
    original_method_name = :"_original_#{method_name}"
    alias_method original_method_name, method_name
    private original_method_name
  end
end

class Baseline
  define_methods_for_testing_memo_wise(self)

  MEMOIZED_METHODS.each do |method_name|
    original_method_name = :"_original_#{method_name}"
    define_method(method_name) do |*args, **kwargs|
      key = [method_name, args, kwargs]
      @cache ||= {}
      @cache.fetch(key) do
        @cache[key] = send(original_method_name, *args, **kwargs)
      end
    end
  end
end

class Initialize
  def initialize
    @cache = {}
  end

  define_methods_for_testing_memo_wise(self)

  MEMOIZED_METHODS.each do |method_name|
    original_method_name = :"_original_#{method_name}"
    define_method(method_name) do |*args, **kwargs|
      key = [method_name, args, kwargs]
      @cache.fetch(key) do
        @cache[key] = send(original_method_name, *args, **kwargs)
      end
    end
  end
end

class ModuleEval
  def initialize
    @cache = {}
  end

  define_methods_for_testing_memo_wise(self)

  (MEMOIZED_METHODS - [:positional_splat_keyword_and_double_splat_args]).each do |method_name|
    module_eval <<-END_OF_METHOD
      def #{method_name}(*args, **kwargs)
        key = [:#{method_name}, args, kwargs]
        @cache.fetch(key) do
          @cache[key] = _original_#{method_name}(*args, **kwargs)
        end
      end
    END_OF_METHOD
  end

  def self.define_positional_splat_keyword_and_double_splat(target)
    target.module_eval <<-END_OF_METHOD
      def positional_splat_keyword_and_double_splat_args(*args, **kwargs)
        key = [:positional_splat_keyword_and_double_splat_args, args, kwargs]
        @cache.fetch(key) do
          @cache[key] = _original_positional_splat_keyword_and_double_splat_args(*args, **kwargs)
        end
      end
    END_OF_METHOD
  end
  define_positional_splat_keyword_and_double_splat(self)
end

class SmallerKeys
  def initialize
    @cache = {}
  end

  REPEATS = ["(a, *args, b:, **kwargs)"]

  define_methods_for_testing_memo_wise(self)

  def self.define_no_args(target)
    NO_ARGS_METHODS.each do |method_name|
      target.module_eval <<-END_OF_METHOD
        def #{method_name}
          @cache.fetch(:#{method_name}) do
            @cache[:#{method_name}] = _original_#{method_name}
          end
        end
      END_OF_METHOD
    end
  end
  define_no_args(self)

  module_eval <<-END_OF_METHOD
    def one_positional_arg(a)
      key = [:one_positional_arg, a]
      @cache.fetch(key) do
        @cache[key] = _original_one_positional_arg(a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_args(a, b)
      key = [:positional_args, a, b]
      @cache.fetch(key) do
        @cache[key] = _original_positional_args(a, b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def one_keyword_arg(a:)
      key = [:one_keyword_arg, a]
      @cache.fetch(key) do
        @cache[key] = _original_one_keyword_arg(a: a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_args(a:, b:)
      key = [:keyword_args, a, b]
      @cache.fetch(key) do
        @cache[key] = _original_keyword_args(a: a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_keyword_args(a, b:)
      key = [:positional_and_keyword_args, a, b]
      @cache.fetch(key) do
        @cache[key] = _original_positional_and_keyword_args(a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_splat_args(*args)
      key = [:positional_and_splat_args, args]
      @cache.fetch(key) do
        @cache[key] = _original_positional_and_splat_args(*args)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_and_double_splat_args(**kwargs)
      key = [:keyword_and_double_splat_args, kwargs]
      @cache.fetch(key) do
        @cache[key] = _original_keyword_and_double_splat_args(**kwargs)
      end
    end
  END_OF_METHOD

  ModuleEval.define_positional_splat_keyword_and_double_splat(self)
end

class TwoLevelHash
  def initialize
    @cache = {}
  end

  REPEATS = ["()"]

  define_methods_for_testing_memo_wise(self)

  SmallerKeys.define_no_args(self)

  module_eval <<-END_OF_METHOD
    def one_positional_arg(a)
      hash = (@cache[:one_positional_arg] ||= {})
      hash.fetch(a) do
        hash[a] = _original_one_positional_arg(a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_args(a, b)
      hash = (@cache[:positional_args] ||= {})
      key = [a, b]
      hash.fetch(key) do
        hash[key] = _original_positional_args(a, b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def one_keyword_arg(a:)
      hash = (@cache[:one_keyword_arg] ||= {})
      hash.fetch(a) do
        hash[a] = _original_one_keyword_arg(a: a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_args(a:, b:)
      hash = (@cache[:keyword_args] ||= {})
      key = [a, b]
      hash.fetch(key) do
        hash[key] = _original_keyword_args(a: a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_keyword_args(a, b:)
      hash = (@cache[:positional_and_keyword_args] ||= {})
      key = [a, b]
      hash.fetch(key) do
        hash[key] = _original_positional_and_keyword_args(a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_splat_args(*args)
      hash = (@cache[:positional_and_splat_args] ||= {})
      hash.fetch(args) do
        hash[args] = _original_positional_and_splat_args(*args)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_and_double_splat_args(**kwargs)
      hash = (@cache[:keyword_and_double_splat_args] ||= {})
      hash.fetch(kwargs) do
        hash[kwargs] = _original_keyword_and_double_splat_args(**kwargs)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      hash = (@cache[:positional_splat_keyword_and_double_splat_args] ||= {})
      key = [args, kwargs]
      hash.fetch(key) do
        hash[key] = _original_positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      end
    end
  END_OF_METHOD
end

class ArrayCache
  def initialize
    @cache = []
    @cache_sentinels = []
  end

  define_methods_for_testing_memo_wise(self)

  def self.define_no_args(target)
    NO_ARGS_METHODS.each_with_index do |method_name, index|
      target.module_eval <<-END_OF_METHOD
        def #{method_name}
          output = @cache[#{index}]
          if output || @cache_sentinels[#{index}]
            output
          else
            @cache_sentinels[#{index}] = true
            @cache[#{index}] = _original_#{method_name}
          end
        end
      END_OF_METHOD
    end
  end
  define_no_args(self)

  module_eval <<-END_OF_METHOD
    def one_positional_arg(a)
      hash = (@cache[2] ||= {})
      hash.fetch(a) do
        hash[a] = _original_one_positional_arg(a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_args(a, b)
      hash = (@cache[3] ||= {})
      key = [a, b]
      hash.fetch(key) do
        hash[key] = _original_positional_args(a, b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def one_keyword_arg(a:)
      hash = (@cache[4] ||= {})
      hash.fetch(a) do
        hash[a] = _original_one_keyword_arg(a: a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_args(a:, b:)
      hash = (@cache[5] ||= {})
      key = [a, b]
      hash.fetch(key) do
        hash[key] = _original_keyword_args(a: a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_keyword_args(a, b:)
      hash = (@cache[6] ||= {})
      key = [a, b]
      hash.fetch(key) do
        hash[key] = _original_positional_and_keyword_args(a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_splat_args(*args)
      hash = (@cache[7] ||= {})
      hash.fetch(args) do
        hash[args] = _original_positional_and_splat_args(*args)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_and_double_splat_args(**kwargs)
      hash = (@cache[8] ||= {})
      hash.fetch(kwargs) do
        hash[kwargs] = _original_keyword_and_double_splat_args(**kwargs)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      hash = (@cache[9] ||= {})
      key = [args, kwargs]
      hash.fetch(key) do
        hash[key] = _original_positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      end
    end
  END_OF_METHOD
end

class Conditional
  def initialize
    @cache = []
    @cache_sentinels = []
  end

  REPEATS = ["()"]

  define_methods_for_testing_memo_wise(self)

  ArrayCache.define_no_args(self)

  module_eval <<-END_OF_METHOD
    def one_positional_arg(a)
      hash = (@cache[2] ||= {})
      if hash.key?(a)
        hash[a]
      else
        hash[a] = _original_one_positional_arg(a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_args(a, b)
      hash = (@cache[3] ||= {})
      key = [a, b]
      if hash.key?(key)
        hash[key]
      else
        hash[key] = _original_positional_args(a, b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def one_keyword_arg(a:)
      hash = (@cache[4] ||= {})
      if hash.key?(a)
        hash[a]
      else
        hash[a] = _original_one_keyword_arg(a: a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_args(a:, b:)
      hash = (@cache[5] ||= {})
      key = [a, b]
      if hash.key?(key)
        hash[key]
      else
        hash[key] = _original_keyword_args(a: a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_keyword_args(a, b:)
      hash = (@cache[6] ||= {})
      key = [a, b]
      if hash.key?(key)
        hash[key]
      else
        hash[key] = _original_positional_and_keyword_args(a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_splat_args(*args)
      hash = (@cache[7] ||= {})
      if hash.key?(args)
        hash[args]
      else
        hash[args] = _original_positional_and_splat_args(*args)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_and_double_splat_args(**kwargs)
      hash = (@cache[8] ||= {})
      if hash.key?(kwargs)
        hash[kwargs]
      else
        hash[kwargs] = _original_keyword_and_double_splat_args(**kwargs)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      hash = (@cache[9] ||= {})
      key = [args, kwargs]
      if hash.key?(key)
        hash[key]
      else
        hash[key] = _original_positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      end
    end
  END_OF_METHOD
end

class TruthyOptimization
  def initialize
    @cache = []
    @cache_sentinels = []
  end

  REPEATS = ["()"]

  define_methods_for_testing_memo_wise(self)

  ArrayCache.define_no_args(self)

  module_eval <<-END_OF_METHOD
    def one_positional_arg(a)
      hash = (@cache[2] ||= {})
      output = hash[a]
      if output || hash.key?(a)
        output
      else
        hash[a] = _original_one_positional_arg(a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_args(a, b)
      hash = (@cache[3] ||= {})
      key = [a, b]
      output = hash[key]
      if output || hash.key?(key)
        output
      else
        hash[key] = _original_positional_args(a, b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def one_keyword_arg(a:)
      hash = (@cache[4] ||= {})
      output = hash[a]
      if output || hash.key?(a)
        output
      else
        hash[a] = _original_one_keyword_arg(a: a)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_args(a:, b:)
      hash = (@cache[5] ||= {})
      key = [a, b]
      output = hash[key]
      if output || hash.key?(key)
        output
      else
        hash[key] = _original_keyword_args(a: a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_keyword_args(a, b:)
      hash = (@cache[6] ||= {})
      key = [a, b]
      output = hash[key]
      if output || hash.key?(key)
        output
      else
        hash[key] = _original_positional_and_keyword_args(a, b: b)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_and_splat_args(*args)
      hash = (@cache[7] ||= {})
      output = hash[args]
      if output || hash.key?(args)
        output
      else
        hash[args] = _original_positional_and_splat_args(*args)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def keyword_and_double_splat_args(**kwargs)
      hash = (@cache[8] ||= {})
      output = hash[kwargs]
      if output || hash.key?(kwargs)
        output
      else
        hash[kwargs] = _original_keyword_and_double_splat_args(**kwargs)
      end
    end
  END_OF_METHOD

  module_eval <<-END_OF_METHOD
    def positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      hash = (@cache[9] ||= {})
      key = [args, kwargs]
      output = hash[key]
      if output || hash.key?(key)
        output
      else
        hash[key] = _original_positional_splat_keyword_and_double_splat_args(*args, **kwargs)
      end
    end
  END_OF_METHOD
end

BENCHMARK_CLASSES = [
  Baseline,
  Initialize,
  ModuleEval,
  SmallerKeys,
  TwoLevelHash,
  ArrayCache,
  Conditional,
  TruthyOptimization
]

class BenchmarkSuiteWithoutGC
  def warming(*) ; run_gc ; end
  def running(*) ; run_gc ; end
  def warmup_stats(*); end
  def add_report(*); end

  private

  def run_gc
    GC.enable
    GC.start
    GC.disable
  end
end
suite = BenchmarkSuiteWithoutGC.new

always_truthy_benchmark_lambdas = [
  lambda do |x, instance|
    instance.no_args
    x.report("#{instance.class}: ()") { instance.no_args }
  end,
  lambda do |x, instance|
    instance.one_positional_arg(1)
    x.report("#{instance.class}: (a)") { instance.one_positional_arg(1) }
  end,
  lambda do |x, instance|
    instance.one_keyword_arg(a: 1)
    x.report("#{instance.class}: (a:)") { instance.one_keyword_arg(a: 1) }
  end,
  lambda do |x, instance|
    instance.positional_args(1, 2)
    x.report("#{instance.class}: (a, b)") { instance.positional_args(1, 2) }
  end,
  lambda do |x, instance|
    instance.keyword_args(a: 1, b: 2)
    x.report("#{instance.class}: (a:, b:)") { instance.keyword_args(a: 1, b: 2) }
  end,
  lambda do |x, instance|
    instance.positional_and_keyword_args(1, b: 2)
    x.report("#{instance.class}: (a, b:)") { instance.positional_and_keyword_args(1, b: 2) }
  end,
  lambda do |x, instance|
    instance.positional_and_splat_args(1, 2)
    x.report("#{instance.class}: (a, *args)") { instance.positional_and_splat_args(1, 2) }
  end,
  lambda do |x, instance|
    instance.keyword_and_double_splat_args(a: 1, b: 2)
    x.report("#{instance.class}: (a:, **kwargs)") { instance.keyword_and_double_splat_args(a: 1, b: 2) }
  end,
  lambda do |x, instance|
    instance.positional_splat_keyword_and_double_splat_args(1, 2, b: 3, a: 4)
    x.report("#{instance.class}: (a, *args, b:, **kwargs)") { instance.positional_splat_keyword_and_double_splat_args(1, 2, b: 3, a: 4) }
  end
]

sometimes_falsey_benchmark_lambdas = [
  lambda do |x, instance|
    instance.no_args_falsey
    instance.no_args

    x.report("#{instance.class}: ()") do
      instance.no_args_falsey
      N_TRUTHY_RESULTS.times { instance.no_args }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, _| instance.one_positional_arg(a) }

    x.report("#{instance.class}: (a)") do
      ARGUMENTS.each { |a, _| instance.one_positional_arg(a) }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, _| instance.one_keyword_arg(a: a) }

    x.report("#{instance.class}: (a:)") do
      ARGUMENTS.each { |a, _| instance.one_keyword_arg(a: a) }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, b| instance.positional_args(a, b) }

    x.report("#{instance.class}: (a, b)") do
      ARGUMENTS.each { |a, b| instance.positional_args(a, b) }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, b| instance.keyword_args(a: a, b: b) }

    x.report("#{instance.class}: (a:, b:)") do
      ARGUMENTS.each { |a, b| instance.keyword_args(a: a, b: b) }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, b| instance.positional_and_keyword_args(a, b: b) }

    x.report("#{instance.class}: (a, b:)") do
      ARGUMENTS.each { |a, b| instance.positional_and_keyword_args(a, b: b) }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, b| instance.positional_and_splat_args(a, b) }

    x.report("#{instance.class}: (a, *args)") do
      ARGUMENTS.each { |a, b| instance.positional_and_splat_args(a, b) }
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each { |a, b| instance.keyword_and_double_splat_args(a: a, b: b) }

    x.report("#{instance.class}: (a:, **kwargs)") do
      ARGUMENTS.each do |a, b|
        instance.keyword_and_double_splat_args(a: a, b: b)
      end
    end
  end,
  lambda do |x, instance|
    ARGUMENTS.each do |a, b|
      instance.positional_splat_keyword_and_double_splat_args(a, b, b: b, a: a)
    end

    x.report("#{instance.class}: (a, *args, b:, **kwargs)") do
      ARGUMENTS.each do |a, b|
        instance.
          positional_splat_keyword_and_double_splat_args(a, b, b: b, a: a)
      end
    end
  end
]

skipped = BENCHMARK_CLASSES.flat_map do |benchmark|
  if defined?(benchmark::REPEATS)
    benchmark.const_get(:REPEATS).map { "#{benchmark}: #{_1}" }
  end
end.compact.to_set

puts "Which benchmarks do you want to run? [T/<n>]"
puts "T: Always truthy values"
puts "<n>: 1/<n> falsey values"

until (value = gets.chomp.upcase) =~ /\A(T|\d+)\z/
  puts "Try again: T/<n>"
end

if value == "T"
  benchmark_lambdas = always_truthy_benchmark_lambdas
  puts "Running benchmarks with always truthy values"
else
  benchmark_lambdas = sometimes_falsey_benchmark_lambdas
  N_UNIQUE_ARGUMENTS = value.to_i # Changing this impacts the frequency of falsey values.
  ARGUMENTS = Array.new(N_UNIQUE_ARGUMENTS) { |i| [i, i + 1] }
  N_TRUTHY_RESULTS = N_UNIQUE_ARGUMENTS - 1
  puts "Running benchmarks with #{100 / N_UNIQUE_ARGUMENTS}% falsey values"
end

benchmark_lambdas.map do |benchmark|
  json_file = Tempfile.new

  Benchmark.ips do |x|
    x.config(suite: suite, warmup: 2, time: 20)

    BENCHMARK_CLASSES.each do |klass|
      instance = klass.new
      benchmark.call(x, instance)
    end

    x.compare!
    x.json! json_file.path
  end

  JSON.parse(json_file.read)
end.each_with_index do |benchmark_json, i|
  # We print a comparison table after we run each benchmark to copy into our
  # README.md

  # MemoWise will not appear in the comparison table, but we will use it to
  # compare against other gems' benchmarks
  baseline = benchmark_json.find { _1["name"].include?("Baseline") }

  # Print headers based on the first benchmark_json
  if i.zero?
    benchmark_headers = benchmark_json.map do |benchmark_gem|
      benchmark_gem["name"].split(":").first
    end.join("|")
    puts "|Method arguments|#{benchmark_headers}|"
    puts "#{'|--' * (benchmark_json.size + 1)}|"
  end

  output_str = benchmark_json.map do |bgem|
    if skipped.include?(bgem["name"])
      "unchanged"
    else
      # "%.2f" % 12.345 => "12.34" (instead of "12.35")
      #   See: https://bugs.ruby-lang.org/issues/12548
      # 1.00.round(2).to_s => "1.0" (instead of "1.00")
      #
      # So to round and format correctly, we first use Float#round and then %
      "%.2fx" % (bgem["central_tendency"] / baseline["central_tendency"]).round(2)
    end
  end.join("|")

  name = baseline["name"].partition(": ").last
  puts "|`#{name}`#{' (none)' if name == '()'}|#{output_str}|"
end