Refactor ActiveRecord::Type::TypeMap

A good part of the complexity was to support the HashLookupTypeMap
subclass that's only used by the Postgres adapter.

In the end they have a similar-ish interface but this inheritance
doesn't help much.

Worse, adapters using `TypeMap` doesn't need extra arguments. By
skipping this unused feature, we can drastically reduce the memory
footprint, as this feature imposed a 320B per entry overhead.

Author: byroot

activerecord/lib/active_record/type/hash_lookup_type_map.rb

@@ -2,7 +2,40 @@
 
 module ActiveRecord
   module Type
-    class HashLookupTypeMap < TypeMap # :nodoc:
+    class HashLookupTypeMap # :nodoc:
+      def initialize(parent = nil)
+        @mapping = {}
+        @cache = Concurrent::Map.new do |h, key|
+          h.fetch_or_store(key, Concurrent::Map.new)
+        end
+      end
+
+      def lookup(lookup_key, *args)
+        fetch(lookup_key, *args) { Type.default_value }
+      end
+
+      def fetch(lookup_key, *args, &block)
+        @cache[lookup_key].fetch_or_store(args) do
+          perform_fetch(lookup_key, *args, &block)
+        end
+      end
+
+      def register_type(key, value = nil, &block)
+        raise ::ArgumentError unless value || block
+
+        if block
+          @mapping[key] = block
+        else
+          @mapping[key] = proc { value }
+        end
+        @cache.clear
+      end
+
+      def clear
+        @mapping.clear
+        @cache.clear
+      end
+
       def alias_type(type, alias_type)
         register_type(type) { |_, *args| lookup(alias_type, *args) }
       end

activerecord/lib/active_record/type/type_map.rb

@@ -8,55 +8,49 @@ class TypeMap # :nodoc:
       def initialize(parent = nil)
         @mapping = {}
         @parent = parent
-        @cache = Concurrent::Map.new do |h, key|
-          h.fetch_or_store(key, Concurrent::Map.new)
-        end
+        @cache = Concurrent::Map.new
       end
 
-      def lookup(lookup_key, *args)
-        fetch(lookup_key, *args) { Type.default_value }
+      def lookup(lookup_key)
+        fetch(lookup_key) { Type.default_value }
       end
 
-      def fetch(lookup_key, *args, &block)
-        @cache[lookup_key].fetch_or_store(args) do
-          perform_fetch(lookup_key, *args, &block)
+      def fetch(lookup_key, &block)
+        @cache.fetch_or_store(lookup_key) do
+          perform_fetch(lookup_key, &block)
         end
       end
 
       def register_type(key, value = nil, &block)
         raise ::ArgumentError unless value || block
-        @cache.clear
 
         if block
           @mapping[key] = block
         else
           @mapping[key] = proc { value }
         end
+        @cache.clear
       end
 
       def alias_type(key, target_key)
-        register_type(key) do |sql_type, *args|
+        register_type(key) do |sql_type|
           metadata = sql_type[/\(.*\)/, 0]
-          lookup("#{target_key}#{metadata}", *args)
+          lookup("#{target_key}#{metadata}")
         end
       end
 
-      def clear
-        @mapping.clear
-      end
-
       protected
-        def perform_fetch(lookup_key, *args, &block)
+        def perform_fetch(lookup_key)
           matching_pair = @mapping.reverse_each.detect do |key, _|
             key === lookup_key
           end
 
           if matching_pair
-            matching_pair.last.call(lookup_key, *args)
+            matching_pair.last.call(lookup_key)
           elsif @parent
-            @parent.perform_fetch(lookup_key, *args, &block)
+            @parent.perform_fetch(lookup_key)
           else
-            yield lookup_key, *args
+            yield lookup_key
           end
         end
     end

activerecord/test/cases/type/type_map_test.rb

@@ -4,16 +4,60 @@
 
 module ActiveRecord
   module Type
-    class TypeMapTest < ActiveRecord::TestCase
+    module TypeMapSharedTests
       def test_default_type
-        mapping = TypeMap.new
+        mapping = klass.new
 
         assert_kind_of Value, mapping.lookup(:undefined)
       end
 
+      def test_requires_value_or_block
+        mapping = klass.new
+
+        assert_raises(ArgumentError) do
+          mapping.register_type(/only key/i)
+        end
+      end
+
+      def test_fetch
+        mapping = klass.new
+        mapping.register_type(1, "string")
+
+        assert_equal "string", mapping.fetch(1) { "int" }
+        assert_equal "int", mapping.fetch(2) { "int" }
+      end
+
+      def test_fetch_memoizes
+        mapping = klass.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_register_clears_cache
+        mapping = klass.new
+
+        mapping.register_type(1, "string")
+        mapping.lookup(1)
+        mapping.register_type(1, "int")
+
+        assert_equal "int", mapping.lookup(1)
+      end
+    end
+
+    class TypeMapTest < ActiveRecord::TestCase
+      include TypeMapSharedTests
+
       def test_registering_types
         boolean = Boolean.new
-        mapping = TypeMap.new
+        mapping = klass.new
 
         mapping.register_type(/boolean/i, boolean)
 
@@ -23,26 +67,17 @@ def test_registering_types
       def test_overriding_registered_types
         time = Time.new
         timestamp = DateTime.new
-        mapping = TypeMap.new
+        mapping = klass.new
 
         mapping.register_type(/time/i, time)
         mapping.register_type(/time/i, timestamp)
 
         assert_equal mapping.lookup("time"), timestamp
       end
 
-      def test_fuzzy_lookup
-        string = +""
-        mapping = TypeMap.new
-
-        mapping.register_type(/varchar/i, string)
-
-        assert_equal mapping.lookup("varchar(20)"), string
-      end
-
       def test_aliasing_types
         string = +""
-        mapping = TypeMap.new
+        mapping = klass.new
 
         mapping.register_type(/string/i, string)
         mapping.alias_type(/varchar/i, "string")
@@ -53,17 +88,17 @@ def test_aliasing_types
       def test_changing_type_changes_aliases
         time = Time.new
         timestamp = DateTime.new
-        mapping = TypeMap.new
+        mapping = klass.new
 
         mapping.register_type(/timestamp/i, time)
         mapping.alias_type(/datetime/i, "timestamp")
         mapping.register_type(/timestamp/i, timestamp)
 
-        assert_equal mapping.lookup("datetime"), timestamp
+        assert_equal timestamp, mapping.lookup("datetime")
       end
 
       def test_aliases_keep_metadata
-        mapping = TypeMap.new
+        mapping = klass.new
 
         mapping.register_type(/decimal/i) { |sql_type| sql_type }
         mapping.alias_type(/number/i, "decimal")
@@ -72,10 +107,19 @@ def test_aliases_keep_metadata
         assert_equal mapping.lookup("number"), "decimal"
       end
 
+      def test_fuzzy_lookup
+        string = +""
+        mapping = klass.new
+
+        mapping.register_type(/varchar/i, string)
+
+        assert_equal mapping.lookup("varchar(20)"), string
+      end
+
       def test_register_proc
         string = +""
         binary = Binary.new
-        mapping = TypeMap.new
+        mapping = klass.new
 
         mapping.register_type(/varchar/i) do |type|
           if type.include?("(")
@@ -89,29 +133,40 @@ def test_register_proc
         assert_equal mapping.lookup("varchar"), binary
       end
 
+      def test_parent_fallback
+        boolean = Boolean.new
+
+        parent = klass.new
+        parent.register_type(/boolean/i, boolean)
+
+        mapping = klass.new(parent)
+        assert_equal boolean, mapping.lookup("boolean")
+      end
+
+      private
+        def klass
+          TypeMap
+        end
+    end
+
+    class HashLookupTypeMapTest < ActiveRecord::TestCase
+      include TypeMapSharedTests
+
       def test_additional_lookup_args
-        mapping = TypeMap.new
+        mapping = HashLookupTypeMap.new
 
-        mapping.register_type(/varchar/i) do |type, limit|
+        mapping.register_type("varchar") do |type, limit|
           if limit > 255
             "text"
           else
             "string"
           end
         end
-        mapping.alias_type(/string/i, "varchar")
+        mapping.alias_type("string", "varchar")
 
-        assert_equal mapping.lookup("varchar", 200), "string"
-        assert_equal mapping.lookup("varchar", 400), "text"
-        assert_equal mapping.lookup("string", 400), "text"
-      end
-
-      def test_requires_value_or_block
-        mapping = TypeMap.new
-
-        assert_raises(ArgumentError) do
-          mapping.register_type(/only key/i)
-        end
+        assert_equal "string", mapping.lookup("varchar", 200)
+        assert_equal "text", mapping.lookup("varchar", 400)
+        assert_equal "text", mapping.lookup("string", 400)
       end
 
       def test_lookup_non_strings
@@ -121,68 +176,31 @@ def test_lookup_non_strings
         mapping.register_type(2, "int")
         mapping.alias_type(3, 1)
 
-        assert_equal mapping.lookup(1), "string"
-        assert_equal mapping.lookup(2), "int"
-        assert_equal mapping.lookup(3), "string"
+        assert_equal "string", mapping.lookup(1)
+        assert_equal "int", mapping.lookup(2)
+        assert_equal "string", mapping.lookup(3)
         assert_kind_of Type::Value, mapping.lookup(4)
       end
 
-      def test_fetch
-        mapping = TypeMap.new
-        mapping.register_type(1, "string")
-
-        assert_equal "string", mapping.fetch(1) { "int" }
-        assert_equal "int", mapping.fetch(2) { "int" }
-      end
-
-      def test_fetch_yields_args
-        mapping = TypeMap.new
-
-        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 = HashLookupTypeMap.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")
+      def test_fetch_yields_args
+        mapping = klass.new
 
-        assert_equal "int", mapping.lookup(1)
+        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_parent_fallback
-        boolean = Boolean.new
-
-        parent = TypeMap.new
-        parent.register_type(/boolean/i, boolean)
-
-        mapping = TypeMap.new(parent)
-        assert_equal mapping.lookup("boolean"), boolean
-      end
+      private
+        def klass
+          HashLookupTypeMap
+        end
     end
   end
 end