From 0367e82ce17302405dff9a0a8e81702db33d83f4 Mon Sep 17 00:00:00 2001
From: Herwin <herwinw@users.noreply.github.com>
Date: Mon, 25 Nov 2024 18:23:41 +0100
Subject: [PATCH] Optimize Enumerable#sum for integer ranges

We can simply shortcut this with the Guass formula for summing integers.
It turns out MRI does this too

Benchmarking this with `(1..1000000).sum`, this takes roughly 1.5
seconds on master (on my machine, release build), and only 0.00003
seconds with this patch. MRI is still roughly 10 times as fast.

Inspired by https://mastodon.social/@riffraff/113503953295615142,
mentioned on https://newsletter.shortruby.com/p/edition-115
MRI does not have a separate `Range#sum` method, so it looks like it is
shortcutted in `Enumerable#range`.
---
 src/enumerable.rb               |  4 ++++
 test/natalie/enumerable_test.rb | 13 +++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/src/enumerable.rb b/src/enumerable.rb
index eeeefa1b7..139765d5a 100644
--- a/src/enumerable.rb
+++ b/src/enumerable.rb
@@ -871,6 +871,10 @@ def sort_by(&block)
 
   def sum(init = 0)
     block_given = block_given?
+    if !block_given && init == 0 && is_a?(Range) && first.is_a?(Integer) && self.last.is_a?(Integer)
+      last = exclude_end? ? self.last - 1 : self.last
+      return (first + last) * (last - first + 1) / 2
+    end
     each do |item|
       if block_given
         init += yield item
diff --git a/test/natalie/enumerable_test.rb b/test/natalie/enumerable_test.rb
index f9421d7af..f399fdbff 100644
--- a/test/natalie/enumerable_test.rb
+++ b/test/natalie/enumerable_test.rb
@@ -57,4 +57,17 @@ def each
       -> { [1].zip [], Object.new, [] }.should raise_error(TypeError)
     end
   end
+
+  describe '#sum' do
+    it 'should return the correct result with optimized Range operations' do
+      (1..1000).sum.should == 500500
+      (1...1000).sum.should == 499500
+      (10..20).sum.should == 165
+    end
+
+    it 'falls back to the default implementation when range is not numeric or a start is given' do
+      ('a'..'c').sum('').should == 'abc'
+      (1..3).sum(10).should == 16
+    end
+  end
 end