Support block yielding for `Enumerable#zip`

2 files changed, 29 insertions, 5 deletions

diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index 6a978b1e3..96249db28 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -742,18 +742,34 @@ module Enumerable
   ##
   #  call-seq:
   #     enum.zip(arg, ...)                  -> an_array_of_array
+  #     enum.zip(arg, ...) { |arr| block }  -> nil
   #
   #  Takes one element from <i>enum</i> and merges corresponding
   #  elements from each <i>args</i>.  This generates a sequence of
   #  <em>n</em>-element arrays, where <em>n</em> is one more than the
   #  count of arguments.  The length of the resulting sequence will be
   #  <code>enum#size</code>.  If the size of any argument is less than
-  #  <code>enum#size</code>, <code>nil</code> values are supplied.
+  #  <code>enum#size</code>, <code>nil</code> values are supplied. If
+  #  a block is given, it is invoked for each output array, otherwise
+  #  an array of arrays is returned.
+  #
+  #     a = [ 4, 5, 6 ]
+  #     b = [ 7, 8, 9 ]
+  #
+  #     a.zip(b)                 #=> [[4, 7], [5, 8], [6, 9]]
+  #     [1, 2, 3].zip(a, b)      #=> [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
+  #     [1, 2].zip(a, b)         #=> [[1, 4, 7], [2, 5, 8]]
+  #     a.zip([1, 2], [8])       #=> [[4, 1, 8], [5, 2, nil], [6, nil, nil]]
+  #
+  #     c = []
+  #     a.zip(b) { |x, y| c << x + y }  #=> nil
+  #     c                               #=> [11, 13, 15]
   #
 
-  def zip(*arg)
-    ary = []
+  def zip(*arg, &block)
+    result = block ? nil : []
     arg = arg.map{|a|a.to_a}
+
     i = 0
     self.each do |*val|
       a = []
@@ -763,10 +779,14 @@ module Enumerable
         a.push(arg[idx][i])
         idx += 1
       end
-      ary.push(a)
       i += 1
+      if result.nil?
+        block.call(a)
+      else
+        result.push(a)
+      end
     end
-    ary
+    result
   end
 
   ##
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 7b9989861..ee4d8b844 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -166,6 +166,10 @@ assert("Enumerable#zip") do
   assert_equal [[1, 4, 7], [2, 5, 8], [3, 6, 9]], [1, 2, 3].zip(a, b)
   assert_equal [[1, 4, 7], [2, 5, 8]], [1, 2].zip(a, b)
   assert_equal [[4, 1, 8], [5, 2, nil], [6, nil, nil]], a.zip([1, 2], [8])
+
+  ret = []
+  assert_equal nil, a.zip([1, 2], [8]) { |i| ret << i }
+  assert_equal [[4, 1, 8], [5, 2, nil], [6, nil, nil]], ret
 end
 
 assert("Enumerable#to_h") do