rubocop · koic · Sep 19, 2020 · Sep 19, 2020
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## master (unreleased)
 
+### Changes
+
+* [#170](https://github.com/rubocop-hq/rubocop-performance/pull/170): Extend `Performance/Sum` to register an offense for `map { ... }.sum`. ([@eugeneius][])
+
 ## 1.8.1 (2020-09-19)
 
 ### Bug fixes

diff --git a/docs/modules/ROOT/pages/cops_performance.adoc b/docs/modules/ROOT/pages/cops_performance.adoc
@@ -1651,11 +1651,12 @@ in some Enumerable object can be replaced by `Enumerable#sum` method.
 [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
 
 # good
 [1, 2, 3].sum
 [1, 2, 3].sum(10)
-[1, 2, 3].sum
+[1, 2, 3].sum { |elem| elem ** 2 }
 ----
 
 === References

diff --git a/lib/rubocop/cop/performance/sum.rb b/lib/rubocop/cop/performance/sum.rb
@@ -12,11 +12,12 @@ module Performance
       #   [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
       #
       #   # good
       #   [1, 2, 3].sum
       #   [1, 2, 3].sum(10)
-      #   [1, 2, 3].sum
+      #   [1, 2, 3].sum { |elem| elem ** 2 }
       #
       class Sum < Base
         include RangeHelp
@@ -28,6 +29,15 @@ class Sum < Base
           (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 ?)
@@ -49,6 +59,16 @@ def on_send(node)
               autocorrect(corrector, init, range)
             end
           end
+
+          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 on_block(node)
@@ -74,10 +94,24 @@ def autocorrect(corrector, init, range)
           corrector.replace(range, replacement)
         end
 
+        def autocorrect_sum_map(corrector, sum, map, init)
+          sum_range = sum.receiver.source_range.end.join(sum.source_range.end)
+          map_range = map.loc.selector
+
+          replacement = build_good_method(init)
+
+          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
@@ -88,6 +122,13 @@ def build_method_message(method, init, operation)
           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)

diff --git a/spec/rubocop/cop/performance/sum_spec.rb b/spec/rubocop/cop/performance/sum_spec.rb
@@ -111,4 +111,51 @@
       RUBY
     end
   end
+
+  %i[map collect].each do |method|
+    it "registers an offense and corrects when using `array.#{method} { |elem| elem ** 2 }.sum`" do
+      expect_offense(<<~RUBY, method: method)
+        array.%{method} { |elem| elem ** 2 }.sum
+              ^{method}^^^^^^^^^^^^^^^^^^^^^^^^^ Use `sum { ... }` instead of `%{method} { ... }.sum`.
+      RUBY
+
+      expect_correction(<<~RUBY)
+        array.sum { |elem| elem ** 2 }
+      RUBY
+    end
+
+    it "registers an offense and corrects when using `array.#{method}(&:count).sum`" do
+      expect_offense(<<~RUBY, method: method)
+        array.%{method}(&:count).sum
+              ^{method}^^^^^^^^^^^^^ Use `sum { ... }` instead of `%{method} { ... }.sum`.
+      RUBY
+
+      expect_correction(<<~RUBY)
+        array.sum(&:count)
+      RUBY
+    end
+
+    it "registers an offense and corrects when using `array.#{method}(&:count).sum(10)`" do
+      expect_offense(<<~RUBY, method: method)
+        array.%{method} { |elem| elem ** 2 }.sum(10)
+              ^{method}^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Use `sum(10) { ... }` instead of `%{method} { ... }.sum(10)`.
+      RUBY
+
+      expect_correction(<<~RUBY)
+        array.sum(10) { |elem| elem ** 2 }
+      RUBY
+    end
+
+    it "does not register an offense when using `array.#{method}(&:count).sum { |elem| elem ** 2 }`" do
+      expect_no_offenses(<<~RUBY)
+        array.#{method}(&:count).sum { |elem| elem ** 2 }
+      RUBY
+    end
+
+    it "does not register an offense when using `array.#{method}(&:count).sum(&:count)`" do
+      expect_no_offenses(<<~RUBY)
+        array.#{method}(&:count).sum(&:count)
+      RUBY
+    end
+  end
 end