Reintroduce cache with tests

3 files changed, 56 insertions, 14 deletions

diff --git a/activerecord/lib/active_record/type/hash_lookup_type_map.rb b/activerecord/lib/active_record/type/hash_lookup_type_map.rb
index 90a8c674cb..82d9327fc0 100644
--- a/activerecord/lib/active_record/type/hash_lookup_type_map.rb
+++ b/activerecord/lib/active_record/type/hash_lookup_type_map.rb
@@ -3,13 +3,15 @@ module ActiveRecord
     class HashLookupTypeMap < TypeMap # :nodoc:
       delegate :key?, to: :@mapping
 
-      def fetch(type, *args, &block)
-        @mapping.fetch(type, block).call(type, *args)
-      end
-
       def alias_type(type, alias_type)
         register_type(type) { |_, *args| lookup(alias_type, *args) }
       end
+
+      private
+
+      def perform_fetch(type, *args, &block)
+        @mapping.fetch(type, block).call(type, *args)
+      end
     end
   end
 end
diff --git a/activerecord/lib/active_record/type/type_map.rb b/activerecord/lib/active_record/type/type_map.rb
index 13d4943861..7c194c0cdf 100644
--- a/activerecord/lib/active_record/type/type_map.rb
+++ b/activerecord/lib/active_record/type/type_map.rb
@@ -3,26 +3,22 @@ module ActiveRecord
     class TypeMap # :nodoc:
       def initialize
         @mapping = {}
+        @cache = Hash.new do |h, key|
+          h[key] = {}
+        end
       end
 
       def lookup(lookup_key, *args)
         fetch(lookup_key, *args) { default_value }
       end
 
-      def fetch(lookup_key, *args)
-        matching_pair = @mapping.reverse_each.detect do |key, _|
-          key === lookup_key
-        end
-
-        if matching_pair
-          matching_pair.last.call(lookup_key, *args)
-        else
-          yield lookup_key, *args
-        end
+      def fetch(lookup_key, *args, &block)
+        @cache[lookup_key][args] ||= perform_fetch(lookup_key, *args, &block)
       end
 
       def register_type(key, value = nil, &block)
         raise ::ArgumentError unless value || block
+        @cache.clear
 
         if block
           @mapping[key] = block
@@ -44,6 +40,18 @@ module ActiveRecord
 
       private
 
+      def perform_fetch(lookup_key, *args)
+        matching_pair = @mapping.reverse_each.detect do |key, _|
+          key === lookup_key
+        end
+
+        if matching_pair
+          matching_pair.last.call(lookup_key, *args)
+        else
+          yield lookup_key, *args
+        end
+      end
+
       def default_value
         @default_value ||= Value.new
       end
diff --git a/activerecord/test/cases/type/type_map_test.rb b/activerecord/test/cases/type/type_map_test.rb
index 1e4eba1961..172c6dfc4c 100644
--- a/activerecord/test/cases/type/type_map_test.rb
+++ b/activerecord/test/cases/type/type_map_test.rb
@@ -139,6 +139,38 @@ module ActiveRecord
         assert_equal "foo-1-2-3", mapping.fetch("foo", 1, 2, 3) { |*args| args.join("-") }
         assert_equal "bar-1-2-3", mapping.fetch("bar", 1, 2, 3) { |*args| args.join("-") }
       end
+
+      def test_fetch_memoizes
+        mapping = TypeMap.new
+
+        looked_up = false
+        mapping.register_type(1) do
+          fail if looked_up
+          looked_up = true
+          "string"
+        end
+
+        assert_equal "string", mapping.fetch(1)
+        assert_equal "string", mapping.fetch(1)
+      end
+
+      def test_fetch_memoizes_on_args
+        mapping = TypeMap.new
+        mapping.register_type("foo") { |*args| args.join("-") }
+
+        assert_equal "foo-1-2-3", mapping.fetch("foo", 1, 2, 3) { |*args| args.join("-") }
+        assert_equal "foo-2-3-4", mapping.fetch("foo", 2, 3, 4) { |*args| args.join("-") }
+      end
+
+      def test_register_clears_cache
+        mapping = TypeMap.new
+
+        mapping.register_type(1, "string")
+        mapping.lookup(1)
+        mapping.register_type(1, "int")
+
+        assert_equal "int", mapping.lookup(1)
+      end
     end
   end
 end