Optimize `Range#count` by using `range_size` if possible

3 files changed, 37 insertions, 7 deletions

diff --git a/benchmark/range_count.yml b/benchmark/range_count.yml
new file mode 100644
index 0000000000..58f53a0236
--- /dev/null
+++ b/benchmark/range_count.yml
@@ -0,0 +1,11 @@
+prelude: |
+  r_1 = 1..1
+  r_1k = 1..1000
+  r_1m = 1..1000000
+  r_str = 'a'..'z'
+
+benchmark:
+  'int 1': r_1.count
+  'int 1K': r_1k.count
+  'int 1M': r_1m.count
+  string: r_str.count
diff --git a/range.c b/range.c
index d0e0c4c1af..e57ffef7e5 100644
--- a/range.c
+++ b/range.c
@@ -607,6 +607,10 @@ double_as_int64(double d)
 static int
 is_integer_p(VALUE v)
 {
+    if (rb_integer_type_p(v)) {
+        return true;
+    }
+
     ID id_integer_p;
     VALUE is_int;
     CONST_ID(id_integer_p, "integer?");
@@ -2166,17 +2170,22 @@ range_count(int argc, VALUE *argv, VALUE range)
          * Infinity. Just let it loop. */
         return rb_call_super(argc, argv);
     }
-    else if (NIL_P(RANGE_END(range))) {
-        /* We are confident that the answer is Infinity. */
-        return DBL2NUM(HUGE_VAL);
-    }
-    else if (NIL_P(RANGE_BEG(range))) {
+
+    VALUE beg = RANGE_BEG(range), end = RANGE_END(range);
+
+    if (NIL_P(beg) || NIL_P(end)) {
         /* We are confident that the answer is Infinity. */
         return DBL2NUM(HUGE_VAL);
     }
-    else {
-        return rb_call_super(argc, argv);
+
+    if (is_integer_p(beg)) {
+        VALUE size = range_size(range);
+        if (!NIL_P(size)) {
+            return size;
+        }
     }
+
+    return rb_call_super(argc, argv);
 }
 
 static bool
diff --git a/test/ruby/test_range.rb b/test/ruby/test_range.rb
index 4d2b6294ed..a09108f806 100644
--- a/test/ruby/test_range.rb
+++ b/test/ruby/test_range.rb
@@ -1075,7 +1075,17 @@ class TestRange < Test::Unit::TestCase
   end
 
   def test_count
+    assert_equal 42, (1..42).count
+    assert_equal 41, (1...42).count
+    assert_equal 0, (42..1).count
+    assert_equal 0, (42...1).count
+    assert_equal 2**100, (1..2**100).count
+    assert_equal 6, (1...6.3).count
+    assert_equal 4, ('a'..'d').count
+    assert_equal 3, ('a'...'d').count
+
     assert_equal(Float::INFINITY, (1..).count)
+    assert_equal(Float::INFINITY, (..1).count)
   end
 
   def test_overlap?