Reduce the amount of work performed when instantiating AR models

We don't know which attributes will or won't be used, and we don't want
to create massive bottlenecks at instantiation. Rather than doing *any*
iteration over types and values, we can lazily instantiate the object.

The lazy attribute hash should not fully implement hash, or subclass
hash at any point in the future. It is not meant to be a replacement,
but instead implement its own interface which happens to overlap.

6 files changed, 79 insertions, 27 deletions

diff --git a/activerecord/lib/active_record/attribute_methods/primary_key.rb b/activerecord/lib/active_record/attribute_methods/primary_key.rb
index 9bd333bbac..104d84a1f8 100644
--- a/activerecord/lib/active_record/attribute_methods/primary_key.rb
+++ b/activerecord/lib/active_record/attribute_methods/primary_key.rb
@@ -120,6 +120,7 @@ module ActiveRecord
         def primary_key=(value)
           @primary_key        = value && value.to_s
           @quoted_primary_key = nil
+          @attributes_builder = nil
         end
       end
     end
diff --git a/activerecord/lib/active_record/attribute_set.rb b/activerecord/lib/active_record/attribute_set.rb
index 98ac63c7e1..21c58cbf1d 100644
--- a/activerecord/lib/active_record/attribute_set.rb
+++ b/activerecord/lib/active_record/attribute_set.rb
@@ -2,8 +2,6 @@ require 'active_record/attribute_set/builder'
 
 module ActiveRecord
   class AttributeSet # :nodoc:
-    delegate :keys, to: :initialized_attributes
-
     def initialize(attributes)
       @attributes = attributes
     end
@@ -25,6 +23,10 @@ module ActiveRecord
       attributes.key?(name) && self[name].initialized?
     end
 
+    def keys
+      attributes.initialized_keys
+    end
+
     def fetch_value(name, &block)
       self[name].value(&block)
     end
@@ -43,7 +45,7 @@ module ActiveRecord
     end
 
     def initialize_dup(_)
-      @attributes = attributes.transform_values(&:dup)
+      @attributes = attributes.dup
       super
     end
 
@@ -58,12 +60,6 @@ module ActiveRecord
       end
     end
 
-    def ensure_initialized(key)
-      unless self[key].initialized?
-        write_from_database(key, nil)
-      end
-    end
-
     protected
 
     attr_reader :attributes
diff --git a/activerecord/lib/active_record/attribute_set/builder.rb b/activerecord/lib/active_record/attribute_set/builder.rb
index d4a787f2fe..3946e02d10 100644
--- a/activerecord/lib/active_record/attribute_set/builder.rb
+++ b/activerecord/lib/active_record/attribute_set/builder.rb
@@ -1,35 +1,83 @@
 module ActiveRecord
   class AttributeSet # :nodoc:
     class Builder # :nodoc:
-      attr_reader :types
+      attr_reader :types, :always_initialized
 
-      def initialize(types)
+      def initialize(types, always_initialized = nil)
         @types = types
+        @always_initialized = always_initialized
       end
 
       def build_from_database(values = {}, additional_types = {})
-        attributes = build_attributes_from_values(values, additional_types)
-        add_uninitialized_attributes(attributes)
+        if always_initialized && !values.key?(always_initialized)
+          values[always_initialized] = nil
+        end
+
+        attributes = LazyAttributeHash.new(types, values, additional_types)
         AttributeSet.new(attributes)
       end
 
       private
+    end
+  end
+
+  class LazyAttributeHash
+    delegate :select, :transform_values, to: :materialize
+    delegate :[], :[]=, :freeze, to: :delegate_hash
+
+    def initialize(types, values, additional_types)
+      @types = types
+      @values = values
+      @additional_types = additional_types
+      @materialized = false
+      @delegate_hash = {}
+      assign_default_proc
+    end
+
+    def key?(key)
+      delegate_hash.key?(key) || values.key?(key) || types.key?(key)
+    end
+
+    def initialized_keys
+      delegate_hash.keys | values.keys
+    end
 
-      def build_attributes_from_values(values, additional_types)
-        values.each_with_object({}) do |(name, value), hash|
-          type = additional_types.fetch(name, types[name])
-          hash[name] = Attribute.from_database(name, value, type)
+    def initialize_dup(_)
+      @delegate_hash = delegate_hash.transform_values(&:dup)
+      assign_default_proc
+      super
+    end
+
+    def initialize_clone(_)
+      @delegate_hash = delegate_hash.clone
+      super
+    end
+
+    protected
+
+    attr_reader :types, :values, :additional_types, :delegate_hash
+
+    private
+
+    def assign_default_proc
+      delegate_hash.default_proc = proc do |hash, name|
+        type = additional_types.fetch(name, types[name])
+
+        if values.key?(name)
+          hash[name] = Attribute.from_database(name, values[name], type)
+        elsif type
+          hash[name] = Attribute.uninitialized(name, type)
         end
       end
+    end
 
-      def add_uninitialized_attributes(attributes)
-        types.each_key do |name|
-          next if attributes.key? name
-          type = types[name]
-          attributes[name] =
-            Attribute.uninitialized(name, type)
-        end
+    def materialize
+      unless @materialized
+        values.each_key { |key| delegate_hash[key] }
+        types.each_key { |key| delegate_hash[key] }
+        @materialized = true
       end
+      delegate_hash
     end
   end
 end
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
index 952aeaa703..89d8932e9e 100644
--- a/activerecord/lib/active_record/core.rb
+++ b/activerecord/lib/active_record/core.rb
@@ -536,8 +536,6 @@ module ActiveRecord
     end
 
     def init_internals
-      @attributes.ensure_initialized(self.class.primary_key)
-
       @aggregation_cache        = {}
       @association_cache        = {}
       @readonly                 = false
diff --git a/activerecord/lib/active_record/model_schema.rb b/activerecord/lib/active_record/model_schema.rb
index a444aac23c..adad7774b9 100644
--- a/activerecord/lib/active_record/model_schema.rb
+++ b/activerecord/lib/active_record/model_schema.rb
@@ -231,7 +231,7 @@ module ActiveRecord
       end
 
       def attributes_builder # :nodoc:
-        @attributes_builder ||= AttributeSet::Builder.new(column_types)
+        @attributes_builder ||= AttributeSet::Builder.new(column_types, primary_key)
       end
 
       def column_types # :nodoc:
diff --git a/activerecord/test/cases/attribute_set_test.rb b/activerecord/test/cases/attribute_set_test.rb
index dc20c3c676..a531be482e 100644
--- a/activerecord/test/cases/attribute_set_test.rb
+++ b/activerecord/test/cases/attribute_set_test.rb
@@ -123,6 +123,15 @@ module ActiveRecord
       assert_nil attributes.fetch_value(:bar)
     end
 
+    test "the primary_key is always initialized" do
+      builder = AttributeSet::Builder.new({ foo: Type::Integer.new }, :foo)
+      attributes = builder.build_from_database
+
+      assert attributes.key?(:foo)
+      assert_equal [:foo], attributes.keys
+      assert attributes[:foo].initialized?
+    end
+
     class MyType
       def type_cast_from_user(value)
         return if value.nil?