* enum.c (rb_enum_values_pack): rename from enum_values_pack, and

remove static.

* enumerator.c (lazy_init_iterator, lazy_init_yielder,
  lazy_select_func, lazy_reject_func, lazy_grep_func): handle
  multiple values correctly.

* enumerator.c (lazy_grep): change the behavior when a block is
  given, to be consistent with Enumerable#grep.

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

ChangeLog

@@ -1,3 +1,15 @@
+Thu Mar 15 23:13:36 2012  Shugo Maeda  <shugo@ruby-lang.org>
+
+	* enum.c (rb_enum_values_pack): rename from enum_values_pack, and
+	  remove static.
+
+	* enumerator.c (lazy_init_iterator, lazy_init_yielder,
+	  lazy_select_func, lazy_reject_func, lazy_grep_func): handle
+	  multiple values correctly.
+
+	* enumerator.c (lazy_grep): change the behavior when a block is
+	  given, to be consistent with Enumerable#grep.
+
 Thu Mar 15 19:12:31 2012  Shugo Maeda  <shugo@ruby-lang.org>
 
 	* enumerator.c (lazy_zip): rescue StopIteration returned by

enum.c

@@ -24,8 +24,8 @@ static ID id_next;
 #define id_cmp  idCmp
 #define id_lshift idLTLT
 
-static VALUE
+VALUE
-enum_values_pack(int argc, VALUE *argv)
+rb_enum_values_pack(int argc, VALUE *argv)
 {
     if (argc == 0) return Qnil;
     if (argc == 1) return argv[0];
@@ -33,7 +33,7 @@ enum_values_pack(int argc, VALUE *argv)
 }
 
 #define ENUM_WANT_SVALUE() do { \
-    i = enum_values_pack(argc, argv); \
+    i = rb_enum_values_pack(argc, argv); \
 } while (0)
 
 #define enum_yield rb_yield_values2
@@ -386,7 +386,7 @@ static VALUE
 collect_all(VALUE i, VALUE ary, int argc, VALUE *argv)
 {
     rb_thread_check_ints();
-    rb_ary_push(ary, enum_values_pack(argc, argv));
+    rb_ary_push(ary, rb_enum_values_pack(argc, argv));
 
     return Qnil;
 }
@@ -967,7 +967,7 @@ static VALUE enum_##name##_func(VALUE result, NODE *memo); \
 static VALUE \
 name##_i(VALUE i, VALUE memo, int argc, VALUE *argv) \
 { \
-    return enum_##name##_func(enum_values_pack(argc, argv), RNODE(memo)); \
+    return enum_##name##_func(rb_enum_values_pack(argc, argv), RNODE(memo)); \
 } \
 \
 static VALUE \
@@ -1622,7 +1622,7 @@ member_i(VALUE iter, VALUE args, int argc, VALUE *argv)
 {
     NODE *memo = RNODE(args);
 
-    if (rb_equal(enum_values_pack(argc, argv), memo->u1.value)) {
+    if (rb_equal(rb_enum_values_pack(argc, argv), memo->u1.value)) {
 	memo->u2.value = Qtrue;
 	rb_iter_break();
     }
@@ -1656,7 +1656,7 @@ each_with_index_i(VALUE i, VALUE memo, int argc, VALUE *argv)
 {
     long n = RNODE(memo)->u3.cnt++;
 
-    return rb_yield_values(2, enum_values_pack(argc, argv), INT2NUM(n));
+    return rb_yield_values(2, rb_enum_values_pack(argc, argv), INT2NUM(n));
 }
 
 /*
@@ -1920,7 +1920,7 @@ zip_ary(VALUE val, NODE *memo, int argc, VALUE *argv)
     int i;
 
     tmp = rb_ary_new2(RARRAY_LEN(args) + 1);
-    rb_ary_store(tmp, 0, enum_values_pack(argc, argv));
+    rb_ary_store(tmp, 0, rb_enum_values_pack(argc, argv));
     for (i=0; i<RARRAY_LEN(args); i++) {
 	VALUE e = RARRAY_PTR(args)[i];
 
@@ -1961,7 +1961,7 @@ zip_i(VALUE val, NODE *memo, int argc, VALUE *argv)
     int i;
 
     tmp = rb_ary_new2(RARRAY_LEN(args) + 1);
-    rb_ary_store(tmp, 0, enum_values_pack(argc, argv));
+    rb_ary_store(tmp, 0, rb_enum_values_pack(argc, argv));
     for (i=0; i<RARRAY_LEN(args); i++) {
 	if (NIL_P(RARRAY_PTR(args)[i])) {
 	    rb_ary_push(tmp, Qnil);
@@ -2049,7 +2049,7 @@ static VALUE
 take_i(VALUE i, VALUE args, int argc, VALUE *argv)
 {
     NODE *memo = RNODE(args);
-    rb_ary_push(memo->u1.value, enum_values_pack(argc, argv));
+    rb_ary_push(memo->u1.value, rb_enum_values_pack(argc, argv));
     if (--memo->u3.cnt == 0) rb_iter_break();
     return Qnil;
 }
@@ -2088,7 +2088,7 @@ static VALUE
 take_while_i(VALUE i, VALUE ary, int argc, VALUE *argv)
 {
     if (!RTEST(enum_yield(argc, argv))) rb_iter_break();
-    rb_ary_push(ary, enum_values_pack(argc, argv));
+    rb_ary_push(ary, rb_enum_values_pack(argc, argv));
     return Qnil;
 }
 
@@ -2123,7 +2123,7 @@ drop_i(VALUE i, VALUE args, int argc, VALUE *argv)
 {
     NODE *memo = RNODE(args);
     if (memo->u3.cnt == 0) {
-	rb_ary_push(memo->u1.value, enum_values_pack(argc, argv));
+	rb_ary_push(memo->u1.value, rb_enum_values_pack(argc, argv));
     }
     else {
 	memo->u3.cnt--;

enumerator.c

@@ -1161,33 +1161,50 @@ generator_each(int argc, VALUE *argv, VALUE obj)
 static VALUE
 lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    VALUE args[2], result;
+    VALUE result;
-    args[0] = m;
+    if (argc == 1) {
-    args[1] = val;
+	VALUE args[2];
-    result = rb_yield_values2(2, args);
+	args[0] = m;
+	args[1] = val;
+	result = rb_yield_values2(2, args);
+    }
+    else {
+	VALUE args;
+	int len = rb_long2int((long)argc + 1);
+
+	args = rb_ary_tmp_new(len);
+	rb_ary_push(args, m);
+	if (argc > 0) {
+	    rb_ary_cat(args, argv, argc);
+	}
+	result = rb_yield_values2(RARRAY_LEN(args), RARRAY_PTR(args));
+	RB_GC_GUARD(args);
+    }
     if (result == Qundef) rb_iter_break();
-    return result;
+    return Qnil;
 }
 
 static VALUE
 lazy_init_yielder(VALUE val, VALUE m, int argc, VALUE *argv)
 {
     VALUE result;
-    result = rb_funcall2(m, id_yield, 1, &val);
+    result = rb_funcall2(m, id_yield, argc, argv);
     if (result == Qundef) rb_iter_break();
-    return result;
+    return Qnil;
 }
 
 static VALUE
 lazy_init_block_i(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    return rb_block_call(m, id_each, argc-1, argv+1, lazy_init_iterator, val);
+    rb_block_call(m, id_each, argc-1, argv+1, lazy_init_iterator, val);
+    return Qnil;
 }
 
 static VALUE
 lazy_init_block(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    return rb_block_call(m, id_each, argc-1, argv+1, lazy_init_yielder, val);
+    rb_block_call(m, id_each, argc-1, argv+1, lazy_init_yielder, val);
+    return Qnil;
 }
 
 static VALUE
@@ -1213,7 +1230,7 @@ lazy_initialize(int argc, VALUE *argv, VALUE self)
     }
     generator = generator_allocate(rb_cGenerator);
     rb_block_call(generator, id_initialize, 0, 0,
-		  (rb_block_given_p() ? lazy_init_block_i: lazy_init_block),
+		  (rb_block_given_p() ? lazy_init_block_i : lazy_init_block),
 		  obj);
     enumerator_init(self, generator, meth, argc - offset, argv + offset);
 
@@ -1260,7 +1277,8 @@ lazy_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
     VALUE result = rb_yield_values2(argc - 1, &argv[1]);
 
-    return rb_funcall(argv[0], id_yield, 1, result);
+    rb_funcall(argv[0], id_yield, 1, result);
+    return Qnil;
 }
 
 static VALUE
@@ -1308,15 +1326,12 @@ lazy_flat_map(VALUE obj)
 static VALUE
 lazy_select_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    VALUE element = argv[1];
+    VALUE element = rb_enum_values_pack(argc - 1, argv + 1);
-    VALUE result = rb_yield_values2(argc - 1, &argv[1]);
 
-    if (RTEST(result)) {
+    if (RTEST(rb_yield(element))) {
 	return rb_funcall(argv[0], id_yield, 1, element);
     }
-    else {
+    return Qnil;
-	return result;
-    }
 }
 
 static VALUE
@@ -1332,15 +1347,12 @@ lazy_select(VALUE obj)
 static VALUE
 lazy_reject_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    VALUE element = argv[1];
+    VALUE element = rb_enum_values_pack(argc - 1, argv + 1);
-    VALUE result = rb_yield_values2(argc - 1, &argv[1]);
 
-    if (!RTEST(result)) {
+    if (!RTEST(rb_yield(element))) {
 	return rb_funcall(argv[0], id_yield, 1, element);
     }
-    else {
+    return Qnil;
-	return result;
-    }
 }
 
 static VALUE
@@ -1356,21 +1368,33 @@ lazy_reject(VALUE obj)
 static VALUE
 lazy_grep_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    VALUE element = argv[1];
+    VALUE i = rb_enum_values_pack(argc - 1, argv + 1);
-    VALUE result = rb_funcall(m, id_eqq, 1, element);
+    VALUE result = rb_funcall(m, id_eqq, 1, i);
 
     if (RTEST(result)) {
-	return rb_funcall(argv[0], id_yield, 1, element);
+	rb_funcall(argv[0], id_yield, 1, i);
     }
-    else {
+    return Qnil;
-	return result;
+}
+
+static VALUE
+lazy_grep_iter(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    VALUE i = rb_enum_values_pack(argc - 1, argv + 1);
+    VALUE result = rb_funcall(m, id_eqq, 1, i);
+
+    if (RTEST(result)) {
+	rb_funcall(argv[0], id_yield, 1, rb_yield(i));
     }
+    return Qnil;
 }
 
 static VALUE
 lazy_grep(VALUE obj, VALUE pattern)
 {
-    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_grep_func, pattern);
+    return rb_block_call(rb_cLazy, id_new, 1, &obj,
+			 rb_block_given_p() ? lazy_grep_iter : lazy_grep_func,
+			 pattern);
 }
 
 static VALUE

include/ruby/intern.h

@@ -198,6 +198,7 @@ VALUE rb_fiber_yield(int argc, VALUE *args);
 VALUE rb_fiber_current(void);
 VALUE rb_fiber_alive_p(VALUE);
 /* enum.c */
+VALUE rb_enum_values_pack(int, VALUE*);
 /* enumerator.c */
 VALUE rb_enumeratorize(VALUE, VALUE, int, VALUE *);
 #define RETURN_ENUMERATOR(obj, argc, argv) do {				\

test/ruby/test_lazy_enumerator.rb

@@ -72,6 +72,18 @@ class TestLazyEnumerator < Test::Unit::TestCase
     assert_equal("word", a.current)
   end
 
+  def test_select_multiple_values
+    e = Enumerator.new { |yielder|
+      for i in 1..5
+        yielder.yield(i, i.to_s)
+      end
+    }
+    assert_equal([[2, "2"], [4, "4"]],
+                 e.select {|x| x[0] % 2 == 0})
+    assert_equal([[2, "2"], [4, "4"]],
+                 e.lazy.select {|x| x[0] % 2 == 0}.force)
+  end
+
   def test_map
     a = Step.new(1..3)
     assert_equal(2, a.map {|x| x * 2}.first)
@@ -112,6 +124,18 @@ class TestLazyEnumerator < Test::Unit::TestCase
     assert_equal(nil, a.current)
   end
 
+  def test_reject_multiple_values
+    e = Enumerator.new { |yielder|
+      for i in 1..5
+        yielder.yield(i, i.to_s)
+      end
+    }
+    assert_equal([[2, "2"], [4, "4"]],
+                 e.reject {|x| x[0] % 2 != 0})
+    assert_equal([[2, "2"], [4, "4"]],
+                 e.lazy.reject {|x| x[0] % 2 != 0}.force)
+  end
+
   def test_grep
     a = Step.new('a'..'f')
     assert_equal('c', a.grep(/c/).first)
@@ -122,6 +146,24 @@ class TestLazyEnumerator < Test::Unit::TestCase
     assert_equal(%w[a e], a.lazy.grep(proc {|x| /[aeiou]/ =~ x}).to_a)
   end
 
+  def test_grep_with_block
+    a = Step.new('a'..'f')
+    assert_equal('C', a.grep(/c/) {|i| i.upcase}.first)
+    assert_equal('C', a.lazy.grep(/c/) {|i| i.upcase}.first)
+  end
+
+  def test_grep_multiple_values
+    e = Enumerator.new { |yielder|
+      3.times { |i|
+        yielder.yield(i, i.to_s)
+      }
+    }
+    assert_equal([[2, "2"]], e.grep(proc {|x| x == [2, "2"]}))
+    assert_equal([[2, "2"]], e.lazy.grep(proc {|x| x == [2, "2"]}).force)
+    assert_equal(["22"],
+                 e.lazy.grep(proc {|x| x == [2, "2"]}, &:join).force)
+  end
+
   def test_zip
     a = Step.new(1..3)
     assert_equal([1, "a"], a.zip("a".."c").first)