Repository: rails/rails
File: activerecord/lib/active_record/associations/belongs_to_association.rb (Ruby)

Code before:
            end
          end
        end

        def update_counters_via_scope(klass, foreign_key, by)
          scope = klass.unscoped.where!(primary_key(klass) => foreign_key)
          scope.update_counters(reflection.counter_cache_column => by, touch: reflection.options[:touch])
        end

        def find_target?
          !loaded? && foreign_key_present? && klass
        end

        def require_counter_update?
          reflection.counter_cache_column && owner.persisted?
        end

        def replace_keys(record, force: false)
          reflection_fk = reflection.foreign_key
          if reflection_fk.is_a?(Array)
            target_key_values = record ? Array(primary_key(record.class)).map { |key| record._read_attribute(key) } : []

            if force || reflection_fk.map { |fk| owner._read_attribute(fk) } != target_key_values
              reflection_fk.each_with_index do |key, index|
                owner[key] = target_key_values[index]
              end
            end
          else
            target_key_value = record ? record._read_attribute(primary_key(record.class)) : nil

            if force || owner._read_attribute(reflection_fk) != target_key_value
              owner[reflection_fk] = target_key_value
            end
          end
        end

        def primary_key(klass)
          reflection.association_primary_key(klass)
        end

        def foreign_key_present?
          Array(reflection.foreign_key).all? { |fk| owner._read_attribute(fk) }
        end

        def invertible_for?(record)
          inverse = inverse_reflection_for(record)
          inverse && (inverse.has_one? || inverse.klass.has_many_inversing)
        end

        def stale_state
          Array(reflection.foreign_key).map do |fk|

Review comment (refactor):
It's a nitpick and I'm ok with not doing it in this PR but this method keeps growing and it started to be really hard to follow it. I wonder if it might be worth looking into refactoring it. 
First I think we should be able to get rid of `if reflection_fk.is_a?(Array)` in favor of treating everything as an `Array()`, even single-column PK/FKs. And then perhaps we could semantically split it into two parts: one part for the "replace" flow and the other one is "nullify" flow where `record` is `nil`.

Again, not a blocking comment, just thinking out loud as we check for the `record` presence multiple time and it feels like we could check it once and as early as possible

Code after:
            end
          end
        end

        def update_counters_via_scope(klass, foreign_key, by)
          scope = klass.unscoped.where!(primary_key(klass) => foreign_key)
          scope.update_counters(reflection.counter_cache_column => by, touch: reflection.options[:touch])
        end

        def find_target?
          !loaded? && foreign_key_present? && klass
        end

        def require_counter_update?
          reflection.counter_cache_column && owner.persisted?
        end

        def replace_keys(record, force: false)
          reflection_fk = reflection.foreign_key
          if reflection_fk.is_a?(Array)
            target_key_values = record ? Array(primary_key(record.class)).map { |key| record._read_attribute(key) } : []

            if force || reflection_fk.map { |fk| owner._read_attribute(fk) } != target_key_values
              owner_pk = Array(owner.class.primary_key)
              reflection_fk.each_with_index do |key, index|
                next if record.nil? && owner_pk.include?(key)
                owner[key] = target_key_values[index]
              end
            end
          else
            target_key_value = record ? record._read_attribute(primary_key(record.class)) : nil

            if force || owner._read_attribute(reflection_fk) != target_key_value
              owner[reflection_fk] = target_key_value
            end
          end
        end

        def primary_key(klass)
          reflection.association_primary_key(klass)
        end

        def foreign_key_present?
          Array(reflection.foreign_key).all? { |fk| owner._read_attribute(fk) }
        end

        def invertible_for?(record)
          inverse = inverse_reflection_for(record)
          inverse && (inverse.has_one? || inverse.klass.has_many_inversing)
        end