1 files changed, 63 insertions, 15 deletions
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