From 7fac529d85818816717c98f40ee4c0b1a79abbcc Mon Sep 17 00:00:00 2001
From: J Smith <dark.panda@gmail.com>
Date: Wed, 28 Nov 2018 11:47:53 -0400
Subject: [PATCH] Allow operator methods to be memoized

Memoizing operators will raise a syntax error due to the operator symbols
being injected into the eval'd code. We can allow operator memoization by
using some mangled names instead. To (hopefully) avoid clashing with any
existing method names, we've added some underscores to the names and
upper-cased them.
---
 lib/memoist.rb       | 36 ++++++++++++++++++++++++++++++++--
 test/memoist_test.rb | 46 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/lib/memoist.rb b/lib/memoist.rb
index c49e102..a9d7483 100644
--- a/lib/memoist.rb
+++ b/lib/memoist.rb
@@ -3,6 +3,34 @@
 require 'memoist/version'
 
 module Memoist
+  OPERATOR_METHOD_NAMES = {
+    '[]'  =>  '__ELEMENT_ACCESS_OP__',
+    '[]=' =>  '__ELEMENT_ASSIGNMENT_OP__',
+    '**'  =>  '__EXPOENTIATION_OP__',
+    '!'   =>  '__NEGATION_BANG_OP__',
+    '~'   =>  '__COMPLEMENT_OP__',
+    '+@'  =>  '__UNARY_PLUS_OP__',
+    '-@'  =>  '__UNARY_MINUS_OP__',
+    '*'   =>  '__MULTIPLICATION_OP__',
+    '/'   =>  '__DIVISION_OP__',
+    '%'   =>  '__MODULO_OP__',
+    '+'   =>  '__PLUS_OP__',
+    '-'   =>  '__MINUS_OP__',
+    '>>'  =>  '__BITWISE_SHIFT_RIGHT_OP__',
+    '<<'  =>  '__BITWISE_SHIFT_LEFT_OP__',
+    '&'   =>  '__BITWISE_AND_OP__',
+    '^'   =>  '__BITWISE_EXCLUSIVE_OR_OP__',
+    '|'   =>  '__BITWISE_OR_OP__',
+    '<='  =>  '__LTE_OP__',
+    '<'   =>  '__LT_OP__',
+    '>'   =>  '__GT_OP__',
+    '>='  =>  '__GTE_OP__',
+    '<=>' =>  '__SPACESHIP_OP__',
+    '=='  =>  '__EQUALITY_OP__',
+    '===' =>  '__TRIPLE_EQUALITY_OP__',
+    '=~'  =>  '__MATCH_OP__',
+  }.freeze
+
   def self.extended(extender)
     Memoist.memoist_eval(extender) do
       unless singleton_class.method_defined?(:memoized_methods)
@@ -13,12 +41,16 @@ def self.memoized_methods
     end
   end
 
+  def self.method_name_for(method_name)
+    OPERATOR_METHOD_NAMES.fetch(method_name.to_s, method_name)
+  end
+
   def self.memoized_ivar_for(method_name, identifier = nil)
-    "@#{memoized_prefix(identifier)}_#{escape_punctuation(method_name)}"
+    "@#{memoized_prefix(identifier)}_#{escape_punctuation(method_name_for(method_name))}"
   end
 
   def self.unmemoized_method_for(method_name, identifier = nil)
-    "#{unmemoized_prefix(identifier)}_#{method_name}".to_sym
+    "#{unmemoized_prefix(identifier)}_#{method_name_for(method_name)}".to_sym
   end
 
   def self.memoized_prefix(identifier = nil)
diff --git a/test/memoist_test.rb b/test/memoist_test.rb
index 433f119..24c107b 100644
--- a/test/memoist_test.rb
+++ b/test/memoist_test.rb
@@ -136,6 +136,40 @@ def name
     end
   end
 
+  class OperatorMethods
+    extend Memoist
+
+    OPERATORS = %w{
+      [] []=
+      **
+      ! ~ +@ -@
+      * / %
+      + -
+      >> <<
+      &
+      ^ |
+      <= < > >=
+      <=>
+      == === =~
+    }.freeze
+
+    attr_reader :counter
+
+    def initialize
+      @counter = CallCounter.new
+    end
+
+    OPERATORS.each do |operator|
+      class_eval(<<-EOS, __FILE__, __LINE__ + 1)
+        def #{operator}(other)
+          @counter.call(#{operator.inspect})
+          'foo'
+        end
+        memoize :#{operator}
+      EOS
+    end
+  end
+
   module Rates
     extend Memoist
 
@@ -560,4 +594,16 @@ def test_private_method_memoization
     assert_equal 'Yes', person.send(:is_developer?)
     assert_equal 1, person.is_developer_calls
   end
+
+  def test_operator_method_names
+    operator_methods = OperatorMethods.new
+
+    OperatorMethods::OPERATORS.each do |operator|
+      3.times { operator_methods.send(operator, 'bar') }
+
+      assert_equal 1, operator_methods.counter.count(operator)
+      assert_equal 'foo', operator_methods.send(operator, 'bar', :reload)
+      assert_equal 2, operator_methods.counter.count(operator)
+    end
+  end
 end