1 files changed, 95 insertions, 54 deletions

diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index 3ba6a71366..5ccaa55a32 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -5,93 +5,134 @@ module ActiveRecord::Associations::Builder
     end
 
     def valid_options
-      super + [:foreign_type, :polymorphic, :touch]
+      super + [:foreign_type, :polymorphic, :touch, :counter_cache]
     end
 
-    def constructable?
-      !options[:polymorphic]
+    def self.valid_dependent_options
+      [:destroy, :delete]
     end
 
-    def build
-      reflection = super
-      add_counter_cache_callbacks(reflection) if options[:counter_cache]
-      add_touch_callbacks(reflection)         if options[:touch]
-      reflection
+    def self.define_callbacks(model, reflection)
+      super
+      add_counter_cache_callbacks(model, reflection) if reflection.options[:counter_cache]
+      add_touch_callbacks(model, reflection)         if reflection.options[:touch]
     end
 
-    def add_counter_cache_callbacks(reflection)
-      cache_column = reflection.counter_cache_column
-      foreign_key = reflection.foreign_key
+    def self.define_accessors(mixin, reflection)
+      super
+      add_counter_cache_methods mixin
+    end
+
+    private
+
+    def self.add_counter_cache_methods(mixin)
+      return if mixin.method_defined? :belongs_to_counter_cache_after_create
 
-      mixin.class_eval <<-CODE, __FILE__, __LINE__ + 1
-        def belongs_to_counter_cache_after_create_for_#{name}
-          record = #{name}
-          record.class.increment_counter(:#{cache_column}, record.id) unless record.nil?
-          @_after_create_counter_called = true
+      mixin.class_eval do
+        def belongs_to_counter_cache_after_create(reflection)
+          if record = send(reflection.name)
+            cache_column = reflection.counter_cache_column
+            record.class.increment_counter(cache_column, record.id)
+            @_after_create_counter_called = true
+          end
         end
 
-        def belongs_to_counter_cache_before_destroy_for_#{name}
-          unless destroyed_by_association && destroyed_by_association.foreign_key.to_sym == #{foreign_key.to_sym.inspect}
-            record = #{name}
-            record.class.decrement_counter(:#{cache_column}, record.id) unless record.nil?
+        def belongs_to_counter_cache_before_destroy(reflection)
+          foreign_key = reflection.foreign_key.to_sym
+          unless destroyed_by_association && destroyed_by_association.foreign_key.to_sym == foreign_key
+            record = send reflection.name
+            if record && !self.destroyed?
+              cache_column = reflection.counter_cache_column
+              record.class.decrement_counter(cache_column, record.id)
+            end
           end
         end
 
-        def belongs_to_counter_cache_after_update_for_#{name}
+        def belongs_to_counter_cache_after_update(reflection)
+          foreign_key  = reflection.foreign_key
+          cache_column = reflection.counter_cache_column
+
           if (@_after_create_counter_called ||= false)
             @_after_create_counter_called = false
-          elsif self.#{foreign_key}_changed? && !new_record? && defined?(#{name.to_s.camelize})
-            model = #{name.to_s.camelize}
-            foreign_key_was = self.#{foreign_key}_was
-            foreign_key = self.#{foreign_key}
+          elsif attribute_changed?(foreign_key) && !new_record? && reflection.constructable?
+            model           = reflection.klass
+            foreign_key_was = attribute_was foreign_key
+            foreign_key     = attribute foreign_key
 
             if foreign_key && model.respond_to?(:increment_counter)
-              model.increment_counter(:#{cache_column}, foreign_key)
+              model.increment_counter(cache_column, foreign_key)
             end
             if foreign_key_was && model.respond_to?(:decrement_counter)
-              model.decrement_counter(:#{cache_column}, foreign_key_was)
+              model.decrement_counter(cache_column, foreign_key_was)
             end
           end
         end
-      CODE
+      end
+    end
+
+    def self.add_counter_cache_callbacks(model, reflection)
+      cache_column = reflection.counter_cache_column
 
-      model.after_create   "belongs_to_counter_cache_after_create_for_#{name}"
-      model.before_destroy "belongs_to_counter_cache_before_destroy_for_#{name}"
-      model.after_update   "belongs_to_counter_cache_after_update_for_#{name}"
+      model.after_create lambda { |record|
+        record.belongs_to_counter_cache_after_create(reflection)
+      }
+
+      model.before_destroy lambda { |record|
+        record.belongs_to_counter_cache_before_destroy(reflection)
+      }
+
+      model.after_update lambda { |record|
+        record.belongs_to_counter_cache_after_update(reflection)
+      }
 
       klass = reflection.class_name.safe_constantize
       klass.attr_readonly cache_column if klass && klass.respond_to?(:attr_readonly)
     end
 
-    def add_touch_callbacks(reflection)
-      mixin.class_eval <<-CODE, __FILE__, __LINE__ + 1
-        def belongs_to_touch_after_save_or_destroy_for_#{name}
-          foreign_key_field = #{reflection.foreign_key.inspect}
-          old_foreign_id    = attribute_was(foreign_key_field)
-
-          if old_foreign_id
-            reflection_klass = #{reflection.klass}
-            old_record       = reflection_klass.find_by(reflection_klass.primary_key => old_foreign_id)
+    def self.touch_record(o, foreign_key, name, touch) # :nodoc:
+      old_foreign_id = o.changed_attributes[foreign_key]
 
-            if old_record
-              old_record.touch #{options[:touch].inspect if options[:touch] != true}
-            end
-          end
+      if old_foreign_id
+        association = o.association(name)
+        reflection = association.reflection
+        if reflection.polymorphic?
+          klass = o.public_send("#{reflection.foreign_type}_was").constantize
+        else
+          klass = association.klass
+        end
+        old_record = klass.find_by(klass.primary_key => old_foreign_id)
 
-          record = #{name}
-          unless record.nil? || record.new_record?
-            record.touch #{options[:touch].inspect if options[:touch] != true}
+        if old_record
+          if touch != true
+            old_record.touch touch
+          else
+            old_record.touch
           end
         end
-      CODE
-
-      model.after_save    "belongs_to_touch_after_save_or_destroy_for_#{name}"
-      model.after_touch   "belongs_to_touch_after_save_or_destroy_for_#{name}"
-      model.after_destroy "belongs_to_touch_after_save_or_destroy_for_#{name}"
+      end
+
+      record = o.send name
+      if record && record.persisted?
+        if touch != true
+          record.touch touch
+        else
+          record.touch
+        end
+      end
     end
 
-    def valid_dependent_options
-      [:destroy, :delete]
+    def self.add_touch_callbacks(model, reflection)
+      foreign_key = reflection.foreign_key
+      n           = reflection.name
+      touch       = reflection.options[:touch]
+
+      callback = lambda { |record|
+        BelongsTo.touch_record(record, foreign_key, n, touch)
+      }
+
+      model.after_save    callback
+      model.after_touch   callback
+      model.after_destroy callback
     end
   end
 end