lib/bullet/counter.rb in flyerhzm-bullet-1.5.3 vs lib/bullet/counter.rb in flyerhzm-bullet-1.5.4

@@ -6,10 +6,12 @@
       def start_request
       end
 
       def end_request
         @@klazz_associations = nil
+        @@possible_objects = nil
+        @@impossible_objects = nil
       end
 
       def need_counter_caches?
         !klazz_associations.empty?
       end
@@ -34,15 +36,31 @@
         [counter_cache_messages(path)]
       end
       
       def add_counter_cache(object, associations)
         klazz = object.class
-        klazz_associations[klazz] ||= []
-        klazz_associations[klazz] << associations
-        klazz_associations[klazz].flatten!
-        klazz_associations[klazz].uniq!
+        if (!possible_objects[klazz].nil? and possible_objects[klazz].include?(object)) and
+           (impossible_objects[klazz].nil? or !impossible_objects[klazz].include?(object))
+          klazz_associations[klazz] ||= []
+          klazz_associations[klazz] << associations
+          unique(klazz_associations[klazz])
+        end
       end
+
+      def add_possible_objects(objects)
+        klazz= objects.first.class
+        possible_objects[klazz] ||= []
+        possible_objects[klazz] << objects
+        unique(possible_objects[klazz])
+      end
+
+      def add_impossible_object(object)
+        klazz = object.class
+        impossible_objects[klazz] ||= []
+        impossible_objects[klazz] << object
+        impossible_objects[klazz].uniq!
+      end
       
       private
         def counter_cache_messages(path = nil)
           messages = []
           klazz_associations.each do |klazz, associations|
@@ -51,16 +69,29 @@
               "  #{klazz} => [#{associations.map(&:inspect).join(', ')}]"
             ]
           end
           messages
         end
+        
+        def unique(array)
+          array.flatten!
+          array.uniq!
+        end
 
         def call_stack_messages
           []
         end
         
         def klazz_associations
           @@klazz_associations ||= {}
+        end
+
+        def possible_objects
+          @@possible_objects ||= {}
+        end
+
+        def impossible_objects
+          @@impossible_objects ||= {}
         end
     end
   end
 end