1 files changed, 67 insertions, 29 deletions
diff --git a/activerecord/lib/active_record/attribute_set/builder.rb b/activerecord/lib/active_record/attribute_set/builder.rb
index 0a62c68bfb..73b77d2eda 100644
--- a/activerecord/lib/active_record/attribute_set/builder.rb
+++ b/activerecord/lib/active_record/attribute_set/builder.rb
@@ -1,50 +1,88 @@
 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 = {})
-        build_from_database_pairs(values.keys, values.values, additional_types)
-      end
+        if always_initialized && !values.key?(always_initialized)
+          values[always_initialized] = nil
+        end
 
-      def build_from_database_pairs(columns, values, additional_types)
-        attributes = build_attributes_from_values(columns, values, additional_types)
-        add_uninitialized_attributes(attributes)
+        attributes = LazyAttributeHash.new(types, values, additional_types)
         AttributeSet.new(attributes)
       end
 
       private
+    end
+  end
 
-      def build_attributes_from_values(columns, values, additional_types)
-        # We are performing manual iteration here as this method is a performance
-        # hotspot
-        hash = {}
-        index = 0
-        length = columns.length
-
-        while index < length
-          name = columns[index]
-          value = values[index]
-          type = additional_types.fetch(name, types[name])
-          hash[name] = Attribute.from_database(name, value, type)
-          index += 1
-        end
+  class LazyAttributeHash
+    delegate :select, :transform_values, to: :materialize
+
+    def initialize(types, values, additional_types)
+      @types = types
+      @values = values
+      @additional_types = additional_types
+      @materialized = false
+      @delegate_hash = {}
+    end
 
-        hash
+    def key?(key)
+      delegate_hash.key?(key) || values.key?(key) || types.key?(key)
+    end
+
+    def [](key)
+      if delegate_hash.key?(key)
+        delegate_hash[key]
+      else
+        assign_default_value(key)
       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 []=(key, value)
+      if frozen?
+        raise RuntimeError, "Can't modify frozen hash"
+      end
+      delegate_hash[key] = value
+    end
+
+    def initialized_keys
+      delegate_hash.keys | values.keys
+    end
+
+    def initialize_dup(_)
+      @delegate_hash = delegate_hash.transform_values(&:dup)
+      super
+    end
+
+    protected
+
+    attr_reader :types, :values, :additional_types, :delegate_hash
+
+    private
+
+    def assign_default_value(name)
+      type = additional_types.fetch(name, types[name])
+
+      if values.key?(name)
+        delegate_hash[name] = Attribute.from_database(name, values[name], type)
+      elsif types.key?(name)
+        delegate_hash[name] = Attribute.uninitialized(name, type)
+      end
+    end
+
+    def materialize
+      unless @materialized
+        values.each_key { |key| self[key] }
+        types.each_key { |key| self[key] }
+        @materialized = true
       end
+      delegate_hash
     end
   end
 end