* array.c (rb_ary_cycle): Support for Array#cycle.size

  [Feature #6636]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37505 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

2 files changed, 23 insertions, 1 deletions

diff --git a/array.c b/array.c
index 03d26afdec..5014452c7f 100644
--- a/array.c
+++ b/array.c
@@ -4229,6 +4229,20 @@ rb_ary_sample(int argc, VALUE *argv, VALUE ary)
     return result;
 }
 
+static VALUE
+rb_ary_cycle_size(VALUE self, VALUE args)
+{
+    long mul;
+    VALUE n = Qnil;
+    if (args && (RARRAY_LEN(args) > 0)) {
+	n = RARRAY_PTR(args)[0];
+    }
+    if (RARRAY_LEN(self) == 0) return INT2FIX(0);
+    if (n == Qnil) return DBL2NUM(INFINITY);
+    mul = NUM2LONG(n);
+    if (mul <= 0) return INT2FIX(0);
+    return rb_funcall(rb_ary_length(self), '*', 1, LONG2FIX(mul));
+}
 
 /*
  *  call-seq:
@@ -4258,7 +4272,7 @@ rb_ary_cycle(int argc, VALUE *argv, VALUE ary)
 
     rb_scan_args(argc, argv, "01", &nv);
 
-    RETURN_ENUMERATOR(ary, argc, argv);
+    RETURN_SIZED_ENUMERATOR(ary, argc, argv, rb_ary_cycle_size);
     if (NIL_P(nv)) {
         n = -1;
     }
diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb
index 93d46a6c23..49d7827d35 100644
--- a/test/ruby/test_enumerator.rb
+++ b/test/ruby/test_enumerator.rb
@@ -456,5 +456,13 @@ class TestEnumerator < Test::Unit::TestCase
       (1..59).to_a.repeated_combination(42).size
       # 1.upto(100).inject(:*) / 1.upto(42).inject(:*) / 1.upto(58).inject(:*)
   end
+
+  def test_size_for_cycle
+    assert_equal Float::INFINITY, [:foo].cycle.size
+    assert_equal 10, [:foo, :bar].cycle(5).size
+    assert_equal 0,  [:foo, :bar].cycle(-10).size
+    assert_equal 0,  [].cycle.size
+    assert_equal 0,  [].cycle(5).size
+  end
 end