sum.rb

# frozen_string_literal: true

module RuboCop
  module Cop
    module Performance
      # This cop identifies places where custom code finding the sum of elements
      # in some Enumerable object can be replaced by `Enumerable#sum` method.
      #
      # @example
      #   # bad
      #   [1, 2, 3].inject(:+)
      #   [1, 2, 3].reduce(10, :+)
      #   [1, 2, 3].inject(&:+)
      #   [1, 2, 3].reduce { |acc, elem| acc + elem }
      #   [1, 2, 3].map { |elem| elem ** 2 }.sum
      #   [1, 2, 3].collect(&:count).sum(10)
      #
      #   # good
      #   [1, 2, 3].sum
      #   [1, 2, 3].sum(10)
      #   [1, 2, 3].sum { |elem| elem ** 2 }
      #   [1, 2, 3].sum(10, &:count)
      #
      class Sum < Base
        include RangeHelp
        extend AutoCorrector

        MSG = 'Use `%<good_method>s` instead of `%<bad_method>s`.'
        MSG_IF_NO_INIT_VALUE =
          'Use `%<good_method>s` instead of `%<bad_method>s`, unless calling `%<bad_method>s` on an empty array.'
        RESTRICT_ON_SEND = %i[inject reduce sum].freeze

        def_node_matcher :sum_candidate?, <<~PATTERN
          (send _ ${:inject :reduce} $_init ? ${(sym :+) (block_pass (sym :+))})
        PATTERN

        def_node_matcher :sum_map_candidate?, <<~PATTERN
          (send
            {
              (block $(send _ {:map :collect}) ...)
              $(send _ {:map :collect} (block_pass _))
            }
          :sum $_init ?)
        PATTERN

        def_node_matcher :sum_with_block_candidate?, <<~PATTERN
          (block
            $(send _ {:inject :reduce} $_init ?)
            (args (arg $_acc) (arg $_elem))
            $send)
        PATTERN

        def_node_matcher :acc_plus_elem?, <<~PATTERN
          (send (lvar %1) :+ (lvar %2))
        PATTERN
        alias elem_plus_acc? acc_plus_elem?

        def on_send(node)
          return if empty_array_literal?(node)

          handle_sum_candidate(node)
          handle_sum_map_candidate(node)
        end

        def on_block(node)
          sum_with_block_candidate?(node) do |send, init, var_acc, var_elem, body|
            if acc_plus_elem?(body, var_acc, var_elem) || elem_plus_acc?(body, var_elem, var_acc)
              range = sum_block_range(send, node)
              message = build_block_message(send, init, var_acc, var_elem, body)

              add_offense(range, message: message) do |corrector|
                autocorrect(corrector, init, range)
              end
            end
          end
        end

        private

        def handle_sum_candidate(node)
          sum_candidate?(node) do |method, init, operation|
            range = sum_method_range(node)
            message = build_method_message(node, method, init, operation)

            add_offense(range, message: message) do |corrector|
              autocorrect(corrector, init, range)
            end
          end
        end

        def handle_sum_map_candidate(node)
          sum_map_candidate?(node) do |map, init|
            next if node.block_literal? || node.block_argument?

            message = build_sum_map_message(map.method_name, init)

            add_offense(sum_map_range(map, node), message: message) do |corrector|
              autocorrect_sum_map(corrector, node, map, init)
            end
          end
        end

        def empty_array_literal?(node)
          receiver = node.children.first
          array_literal?(node) && receiver && receiver.children.empty?
        end

        def array_literal?(node)
          receiver = node.children.first
          receiver&.literal? && receiver&.array_type?
        end

        def autocorrect(corrector, init, range)
          return if init.empty?

          replacement = build_good_method(init)

          corrector.replace(range, replacement)
        end

        def autocorrect_sum_map(corrector, sum, map, init)
          sum_range = method_call_with_args_range(sum)
          map_range = method_call_with_args_range(map)

          block_pass = map.last_argument if map.last_argument&.block_pass_type?
          replacement = build_good_method(init, block_pass)

          corrector.remove(sum_range)
          corrector.replace(map_range, ".#{replacement}")
        end

        def sum_method_range(node)
          range_between(node.loc.selector.begin_pos, node.loc.end.end_pos)
        end

        def sum_map_range(map, sum)
          range_between(map.loc.selector.begin_pos, sum.source_range.end.end_pos)
        end

        def sum_block_range(send, node)
          range_between(send.loc.selector.begin_pos, node.loc.end.end_pos)
        end

        def build_method_message(node, method, init, operation)
          good_method = build_good_method(init)
          bad_method = build_method_bad_method(init, method, operation)
          msg = if init.empty? && !array_literal?(node)
                  MSG_IF_NO_INIT_VALUE
                else
                  MSG
                end
          format(msg, good_method: good_method, bad_method: bad_method)
        end

        def build_sum_map_message(method, init)
          sum_method = build_good_method(init)
          good_method = "#{sum_method} { ... }"
          bad_method = "#{method} { ... }.#{sum_method}"
          format(MSG, good_method: good_method, bad_method: bad_method)
        end

        def build_block_message(send, init, var_acc, var_elem, body)
          good_method = build_good_method(init)
          bad_method = build_block_bad_method(send.method_name, init, var_acc, var_elem, body)
          format(MSG, good_method: good_method, bad_method: bad_method)
        end

        def build_good_method(init, block_pass = nil)
          good_method = 'sum'

          args = []
          unless init.empty?
            init = init.first
            args << init.source unless init.int_type? && init.value.zero?
          end
          args << block_pass.source if block_pass
          good_method += "(#{args.join(', ')})" unless args.empty?
          good_method
        end

        def build_method_bad_method(init, method, operation)
          bad_method = "#{method}("
          unless init.empty?
            init = init.first
            bad_method += "#{init.source}, "
          end
          bad_method += if operation.block_pass_type?
                          '&:+)'
                        else
                          ':+)'
                        end
          bad_method
        end

        def build_block_bad_method(method, init, var_acc, var_elem, body)
          bad_method = method.to_s

          unless init.empty?
            init = init.first
            bad_method += "(#{init.source})"
          end
          bad_method += " { |#{var_acc}, #{var_elem}| #{body.source} }"
          bad_method
        end

        def method_call_with_args_range(node)
          node.receiver.source_range.end.join(node.source_range.end)
        end
      end
    end
  end
end