1 files changed, 43 insertions, 59 deletions

diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 4558872a2b..c263edd2c6 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -1,41 +1,17 @@
 module ActiveRecord
   # = Active Record Belongs To Associations
   module Associations
-    class BelongsToAssociation < AssociationProxy #:nodoc:
-      def create(attributes = {})
-        replace(@reflection.create_association(attributes))
-      end
-
-      def build(attributes = {})
-        replace(@reflection.build_association(attributes))
-      end
-
+    class BelongsToAssociation < SingularAssociation #:nodoc:
       def replace(record)
-        counter_cache_name = @reflection.counter_cache_column
-
-        if record.nil?
-          if counter_cache_name && !@owner.new_record?
-            @reflection.klass.decrement_counter(counter_cache_name, previous_record_id) if @owner[@reflection.primary_key_name]
-          end
-
-          @target = @owner[@reflection.primary_key_name] = nil
-        else
-          raise_on_type_mismatch(record)
+        raise_on_type_mismatch(record) if record
 
-          if counter_cache_name && !@owner.new_record? && record.id != @owner[@reflection.primary_key_name]
-            @reflection.klass.increment_counter(counter_cache_name, record.id)
-            @reflection.klass.decrement_counter(counter_cache_name, @owner[@reflection.primary_key_name]) if @owner[@reflection.primary_key_name]
-          end
-
-          @target = (AssociationProxy === record ? record.target : record)
-          @owner[@reflection.primary_key_name] = record_id(record) unless record.new_record?
-          @updated = true
-        end
+        update_counters(record)
+        replace_keys(record)
+        set_inverse_instance(record)
 
-        set_inverse_instance(record, @owner)
+        @updated = true if record
 
-        loaded
-        record
+        self.target = record
       end
 
       def updated?
@@ -43,44 +19,52 @@ module ActiveRecord
       end
 
       private
-        def find_target
-          find_method = if @reflection.options[:primary_key]
-                          "find_by_#{@reflection.options[:primary_key]}"
-                        else
-                          "find"
-                        end
-          the_target = @reflection.klass.send(find_method,
-            @owner[@reflection.primary_key_name],
-            :select     => @reflection.options[:select],
-            :conditions => conditions,
-            :include    => @reflection.options[:include],
-            :readonly   => @reflection.options[:readonly]
-          ) if @owner[@reflection.primary_key_name]
-          set_inverse_instance(the_target, @owner)
-          the_target
+
+        def update_counters(record)
+          counter_cache_name = reflection.counter_cache_column
+
+          if counter_cache_name && owner.persisted? && different_target?(record)
+            if record
+              record.class.increment_counter(counter_cache_name, record.id)
+            end
+
+            if foreign_key_present?
+              klass.decrement_counter(counter_cache_name, target_id)
+            end
+          end
         end
 
-        def foreign_key_present
-          !@owner[@reflection.primary_key_name].nil?
+        # Checks whether record is different to the current target, without loading it
+        def different_target?(record)
+          record.nil? && owner[reflection.foreign_key] ||
+          record.id   != owner[reflection.foreign_key]
+        end
+
+        def replace_keys(record)
+          owner[reflection.foreign_key] = record && record[reflection.association_primary_key]
+        end
+
+        def foreign_key_present?
+          owner[reflection.foreign_key]
         end
 
         # NOTE - for now, we're only supporting inverse setting from belongs_to back onto
         # has_one associations.
-        def we_can_set_the_inverse_on_this?(record)
-          @reflection.has_inverse? && @reflection.inverse_of.macro == :has_one
+        def invertible_for?(record)
+          inverse = inverse_reflection_for(record)
+          inverse && inverse.macro == :has_one
         end
 
-        def record_id(record)
-          record.send(@reflection.options[:primary_key] || :id)
+        def target_id
+          if options[:primary_key]
+            owner.send(reflection.name).try(:id)
+          else
+            owner[reflection.foreign_key]
+          end
         end
 
-        def previous_record_id
-          @previous_record_id ||= if @reflection.options[:primary_key]
-                                    previous_record = @owner.send(@reflection.name)
-                                    previous_record.nil? ? nil : previous_record.id
-                                  else
-                                    @owner[@reflection.primary_key_name]
-                                  end
+        def stale_state
+          owner[reflection.foreign_key].to_s
         end
     end
   end