From 8f57f8f2df6658493d7b8d0a91411b2c0506d2c8 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 54b67e6..e048432 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 49a09d0..6dbe40a 100644
--- a/test/memoist_test.rb
+++ b/test/memoist_test.rb
@@ -130,6 +130,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
 
@@ -547,4 +581,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