diff options
| -rw-r--r-- | ChangeLog | 12 | ||||
| -rw-r--r-- | enum.c | 24 | ||||
| -rw-r--r-- | enumerator.c | 82 | ||||
| -rw-r--r-- | include/ruby/intern.h | 1 | ||||
| -rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 42 |
5 files changed, 120 insertions, 41 deletions
@@ -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 @@ -24,8 +24,8 @@ static ID id_next; #define id_cmp idCmp #define id_lshift idLTLT -static VALUE -enum_values_pack(int argc, VALUE *argv) +VALUE +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--; diff --git a/enumerator.c b/enumerator.c index 09543940f4..27a23998b0 100644 --- a/enumerator.c +++ b/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; - args[0] = m; - args[1] = val; - result = rb_yield_values2(2, args); + VALUE result; + if (argc == 1) { + VALUE args[2]; + 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 result = rb_yield_values2(argc - 1, &argv[1]); + VALUE element = rb_enum_values_pack(argc - 1, argv + 1); - if (RTEST(result)) { + if (RTEST(rb_yield(element))) { return rb_funcall(argv[0], id_yield, 1, element); } - else { - return result; - } + return Qnil; } 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 result = rb_yield_values2(argc - 1, &argv[1]); + VALUE element = rb_enum_values_pack(argc - 1, argv + 1); - if (!RTEST(result)) { + if (!RTEST(rb_yield(element))) { return rb_funcall(argv[0], id_yield, 1, element); } - else { - return result; - } + return Qnil; } 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 result = rb_funcall(m, id_eqq, 1, element); + VALUE i = rb_enum_values_pack(argc - 1, argv + 1); + 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 result; + return Qnil; +} + +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 diff --git a/include/ruby/intern.h b/include/ruby/intern.h index cf36aa6265..e07f7fd781 100644 --- a/include/ruby/intern.h +++ b/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 { \ diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 7064f3b178..b096821191 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/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) |