module Mongoid::Association::Referenced::CounterCache

def self.define_callbacks!(association)
  name = association.name
  cache_column = association.counter_cache_column_name.to_sym
  association.inverse_class.tap do |klass|
    klass.after_update do
      foreign_key = association.foreign_key
      if send("#{foreign_key}_previously_changed?")
        original, current = send("#{foreign_key}_previous_change")
        unless original.nil?
          association.klass.with(persistence_context.for_child(association.klass)) do |_class|
            _class.decrement_counter(cache_column, original)
          end
        end
        if record = __send__(name)
          unless current.nil?
            record[cache_column] = (record[cache_column] || 0) + 1
            record.class.with(record.persistence_context) do |_class|
              _class.increment_counter(cache_column, current) if record.persisted?
            end
          end
        end
      end
    end
    klass.after_create do
      if record = __send__(name)
        record[cache_column] = (record[cache_column] || 0) + 1
        if record.persisted?
          record.class.with(record.persistence_context) do |_class|
            _class.increment_counter(cache_column, record._id)
          end
          record.remove_change(cache_column)
        end
      end
    end
    klass.before_destroy do
      if record = __send__(name)
        record[cache_column] = (record[cache_column] || 0) - 1 unless record.frozen?
        if record.persisted?
          record.class.with(record.persistence_context) do |_class|
            _class.decrement_counter(cache_column, record._id)
          end
          record.remove_change(cache_column)
        end
      end
    end
  end
end

def reset_counters(*counters)
  self.class.with(persistence_context) do |_class|
    _class.reset_counters(self, *counters)
  end
end