Add Enumerable#cycle

2 files changed, 41 insertions, 0 deletions

diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index 0ce1d7605..dcfdb1432 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -523,4 +523,38 @@ module Enumerable
     ary.reverse_each(&block)
     self
   end
+
+  ##
+  #  call-seq:
+  #     enum.cycle(n=nil) { |obj| block }  ->  nil
+  #     enum.cycle(n=nil)                  ->  an_enumerator
+  #
+  #  Calls <i>block</i> for each element of <i>enum</i> repeatedly _n_
+  #  times or forever if none or +nil+ is given.  If a non-positive
+  #  number is given or the collection is empty, does nothing.  Returns
+  #  +nil+ if the loop has finished without getting interrupted.
+  #
+  #  Enumerable#cycle saves elements in an internal array so changes
+  #  to <i>enum</i> after the first pass have no effect.
+  #
+  #  If no block is given, an enumerator is returned instead.
+  #
+  #     a = ["a", "b", "c"]
+  #     a.cycle { |x| puts x }  # print, a, b, c, a, b, c,.. forever.
+  #     a.cycle(2) { |x| puts x }  # print, a, b, c, a, b, c.
+  #
+
+  def cycle(n=nil, &block)
+    if n == nil
+      loop {self.each {|val| block.call(val) } }
+    else
+      raise TypeError, "expected Integer for 1st argument" unless n.kind_of? Integer
+
+      count = 0
+      while count < n
+        self.each {|val| block.call(val) }
+        count += 1
+      end
+    end
+  end
 end
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 127b263fb..68f2781b7 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -122,3 +122,10 @@ assert("Enumerable#reverse_each") do
   assert_equal (1..3), r.reverse_each { |v| a << v }
   assert_equal [3, 2, 1], a
 end
+
+assert("Enumerable#cycle") do
+  a = []
+  ["a", "b", "c"].cycle(2) { |v| a << v }
+  assert_equal ["a", "b", "c", "a", "b", "c"], a
+  assert_raise(TypeError) { ["a", "b", "c"].cycle("a") { |v| a << v } }
+end