Merge pull request #3987 from ksss/enum-lazy-zip

Reimplement `Enumerable#zip` with Enumerator

2 files changed, 42 insertions, 17 deletions

diff --git a/mrbgems/mruby-enumerator/mrblib/enumerator.rb b/mrbgems/mruby-enumerator/mrblib/enumerator.rb
index 1e77af369..7ca1d5eb6 100644
--- a/mrbgems/mruby-enumerator/mrblib/enumerator.rb
+++ b/mrbgems/mruby-enumerator/mrblib/enumerator.rb
@@ -621,25 +621,40 @@ end
 
 module Enumerable
   # use Enumerator to use infinite sequence
-  def zip(*arg)
-    ary = []
-    arg = arg.map{|a|a.each}
-    i = 0
-    self.each do |*val|
-      a = []
-      a.push(val.__svalue)
-      idx = 0
-      while idx < arg.size
-        begin
-          a.push(arg[idx].next)
-        rescue StopIteration
-          a.push(nil)
+  def zip(*args, &block)
+    args = args.map do |a|
+      if a.respond_to?(:to_ary)
+        a.to_ary.to_enum(:each)
+      elsif a.respond_to?(:each)
+        a.to_enum(:each)
+      else
+        raise TypeError, "wrong argument type #{a.class} (must respond to :each)"
+      end
+    end
+
+    result = block ? nil : []
+
+    each do |*val|
+      tmp = [val.__svalue]
+      args.each do |arg|
+        v = if arg.nil?
+          nil
+        else
+          begin
+            arg.next
+          rescue StopIteration
+            nil
+          end
         end
-        idx += 1
+        tmp.push(v)
+      end
+      if result.nil?
+        block.call(tmp)
+      else
+        result.push(tmp)
       end
-      ary.push(a)
-      i += 1
     end
-    ary
+
+    result
   end
 end
diff --git a/mrbgems/mruby-enumerator/test/enumerator.rb b/mrbgems/mruby-enumerator/test/enumerator.rb
index 4c904a81d..428ea0307 100644
--- a/mrbgems/mruby-enumerator/test/enumerator.rb
+++ b/mrbgems/mruby-enumerator/test/enumerator.rb
@@ -544,3 +544,13 @@ assert 'Range#each' do
   end
   assert_equal [1,2,3,4,5], c
 end
+
+assert 'Enumerable#zip' do
+  assert_equal [[1, 10], [2, 11], [3, 12]], [1,2,3].zip(10..Float::INFINITY)
+
+  ret = []
+  assert_equal nil, [1,2,3].zip(10..Float::INFINITY) { |i| ret << i }
+  assert_equal [[1, 10], [2, 11], [3, 12]], ret
+
+  assert_raise(TypeError) { [1].zip(1) }
+end