1 files changed, 34 insertions, 69 deletions

diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 98c1c13524..e80b945dda 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -1,84 +1,50 @@
 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.persisted?
-            @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)
-
-          if counter_cache_name && @owner.persisted? && 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) if record.persisted?
-          @updated = true
-        end
+        record = check_record(record)
 
+        update_counters(record)
+        replace_keys(record)
         set_inverse_instance(record)
 
-        loaded
-        record
+        @updated = true if record
+
+        self.target = record
       end
 
       def updated?
         @updated
       end
 
-      def stale_target?
-        if @target && @target.persisted?
-          target_id   = @target.send(@reflection.association_primary_key).to_s
-          foreign_key = @owner.send(@reflection.primary_key_name).to_s
-
-          target_id != foreign_key
-        else
-          false
-        end
-      end
-
       private
-        def find_target
-          find_method = if @reflection.options[:primary_key]
-                          "find_by_#{@reflection.options[:primary_key]}"
-                        else
-                          "find"
-                        end
+        def update_counters(record)
+          counter_cache_name = @reflection.counter_cache_column
 
-          options = @reflection.options.dup.slice(:select, :include, :readonly)
+          if counter_cache_name && @owner.persisted? && different_target?(record)
+            if record
+              record.class.increment_counter(counter_cache_name, record.id)
+            end
 
-          the_target = with_scope(:find => @scope[:find]) do
-            @reflection.klass.send(find_method,
-              @owner[@reflection.primary_key_name],
-              options
-            ) if @owner[@reflection.primary_key_name]
+            if foreign_key_present?
+              target_klass.decrement_counter(counter_cache_name, target_id)
+            end
           end
-          set_inverse_instance(the_target)
-          the_target
         end
 
-        def construct_find_scope
-          { :conditions => conditions }
+        # 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.primary_key_name].nil?
+        def foreign_key_present?
+          @owner[@reflection.foreign_key]
         end
 
         # NOTE - for now, we're only supporting inverse setting from belongs_to back onto
@@ -88,17 +54,16 @@ module ActiveRecord
           inverse && inverse.macro == :has_one
         end
 
-        def record_id(record)
-          record.send(@reflection.options[:primary_key] || :id)
+        def target_id
+          if @reflection.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