marcandre	2013-01-24 16:05:42 +0900 (Thu, 24 Jan 2013)

  New Revision: 38925

  http://svn.ruby-lang.org/cgi-bin/viewvc.cgi?view=rev&revision=38925

  Log:
    * enumerator.c: Optimize Lazy#zip when passed only arrays
      [Bug #7706]

  Modified files:
    trunk/ChangeLog
    trunk/enumerator.c
    trunk/test/ruby/test_lazy_enumerator.rb

Index: ChangeLog
===================================================================
--- ChangeLog	(revision 38924)
+++ ChangeLog	(revision 38925)
@@ -1,3 +1,8 @@ https://github.com/ruby/ruby/blob/trunk/ChangeLog#L1
+Thu Jan 24 16:05:08 2013  Marc-Andre Lafortune  <ruby-core@m...>
+
+	* enumerator.c: Optimize Lazy#zip when passed only arrays
+	  [Bug #7706]
+
 Thu Jan 24 15:21:17 2013  Marc-Andre Lafortune  <ruby-core@m...>
 
 	* enumerator.c: Fix state handling for Lazy#zip,{drop_take}{_while}
Index: enumerator.c
===================================================================
--- enumerator.c	(revision 38924)
+++ enumerator.c	(revision 38925)
@@ -1602,6 +1602,26 @@ next_stopped(VALUE obj) https://github.com/ruby/ruby/blob/trunk/enumerator.c#L1602
 }
 
 static VALUE
+lazy_zip_arrays_func(VALUE val, VALUE arrays, int argc, VALUE *argv)
+{
+    VALUE yielder, ary, memo;
+    long i, count;
+
+    yielder = argv[0];
+    memo = rb_ivar_get(yielder, id_memo);
+    count = NIL_P(memo) ? 0 : NUM2LONG(memo);
+
+    ary = rb_ary_new2(RARRAY_LEN(arrays) + 1);
+    rb_ary_push(ary, argv[1]);
+    for (i = 0; i < RARRAY_LEN(arrays); i++) {
+	rb_ary_push(ary, rb_ary_entry(RARRAY_PTR(arrays)[i], count));
+    }
+    rb_funcall(yielder, id_yield, 1, ary);
+    rb_ivar_set(yielder, id_memo, LONG2NUM(++count));
+    return Qnil;
+}
+
+static VALUE
 lazy_zip_func(VALUE val, VALUE zip_args, int argc, VALUE *argv)
 {
     VALUE yielder, ary, arg, v;
@@ -1631,15 +1651,27 @@ lazy_zip_func(VALUE val, VALUE zip_args, https://github.com/ruby/ruby/blob/trunk/enumerator.c#L1651
 static VALUE
 lazy_zip(int argc, VALUE *argv, VALUE obj)
 {
-    VALUE ary;
+    VALUE ary, v;
+    long i;
+    rb_block_call_func *func = lazy_zip_arrays_func;
 
     if (rb_block_given_p()) {
 	return rb_call_super(argc, argv);
     }
-    ary = rb_ary_new4(argc, argv);
+
+    ary = rb_ary_new2(argc);
+    for (i = 0; i < argc; i++) {
+	v = rb_check_array_type(argv[i]);
+	if (NIL_P(v)) {
+	    ary = rb_ary_new4(argc, argv);
+	    func = lazy_zip_func;
+	    break;
+	}
+	rb_ary_push(ary, v);
+    }
 
     return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
-					 lazy_zip_func, ary),
+					 func, ary),
 			   ary, lazy_receiver_size);
 }
 
Index: test/ruby/test_lazy_enumerator.rb
===================================================================
--- test/ruby/test_lazy_enumerator.rb	(revision 38924)
+++ test/ruby/test_lazy_enumerator.rb	(revision 38925)
@@ -272,6 +272,20 @@ class TestLazyEnumerator < Test::Unit::T https://github.com/ruby/ruby/blob/trunk/test/ruby/test_lazy_enumerator.rb#L272
     assert_equal([[1, 'a', 'a'], [2, 'b', 'b'], [3, 'c', 'c']]*3, zip.flat_map{zip}.force, bug7696)
   end
 
+  def test_zip_lazy_on_args
+    zip = Step.new(1..2).lazy.zip(42..Float::INFINITY)
+    assert_equal [[1, 42], [2, 43]], zip.force
+  end
+
+  def test_zip_efficient_on_array_args
+    ary = [42, :foo]
+    %i[to_enum enum_for lazy each].each do |forbid|
+      ary.define_singleton_method(forbid){ fail "#{forbid} was called"}
+    end
+    zip = Step.new(1..2).lazy.zip(ary)
+    assert_equal [[1, 42], [2, :foo]], zip.force
+  end
+
   def test_take_rewound
     bug7696 = '[ruby-core:51470]'
     e=(1..42).lazy.take(2)

--
ML: ruby-changes@q...
Info: http://www.atdot.net/~ko1/quickml/