diff options
| -rw-r--r-- | ChangeLog | 14 | ||||
| -rw-r--r-- | enumerator.c | 81 | ||||
| -rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 44 |
3 files changed, 130 insertions, 9 deletions
@@ -1,3 +1,17 @@ +Wed Mar 14 22:01:06 2012 Shugo Maeda <shugo@ruby-lang.org> + + * enumerator.c (lazy_init_iterator): break when Qundef is returned + to make obj.drop(3).take(2) work properly. + + * enumerator.c (lazy_take_while): add Enumerable::Lazy#take_while. + + * enumerator.c (lazy_drop): add Enumerable::Lazy#drop. + + * enumerator.c (lazy_drop_while): add Enumerable::Lazy#drop_while. + + * enumerator.c (InitVM_Enumerator): add Enumerable::Lazy#force as an + alias of to_a. + Wed Mar 14 19:28:40 2012 Shugo Maeda <shugo@ruby-lang.org> * enumerator.c (lazy_take): add Enumerable::Lazy#take. diff --git a/enumerator.c b/enumerator.c index 45fa4f89e6..ab0c39e6c1 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1161,10 +1161,12 @@ generator_each(int argc, VALUE *argv, VALUE obj) static VALUE lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv) { - VALUE args[2]; + VALUE args[2], result; args[0] = m; args[1] = val; - return rb_yield_values2(2, args); + result = rb_yield_values2(2, args); + if (result == Qundef) rb_iter_break(); + return result; } static VALUE @@ -1407,9 +1409,7 @@ lazy_take_func(VALUE val, VALUE args, int argc, VALUE *argv) { NODE *memo = RNODE(args); - if (memo->u3.cnt == 0) { - return Qundef; - } + if (memo->u3.cnt == 0) return Qundef; rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); memo->u3.cnt--; return Qnil; @@ -1430,6 +1430,73 @@ lazy_take(VALUE obj, VALUE n) } static VALUE +lazy_take_while_func(VALUE val, VALUE args, int argc, VALUE *argv) +{ + VALUE result = rb_yield_values2(argc - 1, &argv[1]); + if (!RTEST(result)) return Qundef; + rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); + return Qnil; +} + +static VALUE +lazy_take_while(VALUE obj) +{ + return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_take_while_func, 0); +} + +static VALUE +lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv) +{ + NODE *memo = RNODE(args); + + if (memo->u3.cnt == 0) { + rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); + } + else { + memo->u3.cnt--; + } + return Qnil; +} + +static VALUE +lazy_drop(VALUE obj, VALUE n) +{ + NODE *memo; + long len = NUM2LONG(n); + + if (len < 0) { + rb_raise(rb_eArgError, "attempt to drop negative size"); + } + memo = NEW_MEMO(0, 0, len); + return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_drop_func, + (VALUE) memo); +} + +static VALUE +lazy_drop_while_func(VALUE val, VALUE args, int argc, VALUE *argv) +{ + NODE *memo = RNODE(args); + + if (!memo->u3.state && !RTEST(rb_yield_values2(argc - 1, &argv[1]))) { + memo->u3.state = TRUE; + } + if (memo->u3.state) { + rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); + } + return Qnil; +} + +static VALUE +lazy_drop_while(VALUE obj) +{ + NODE *memo; + + memo = NEW_MEMO(0, 0, FALSE); + return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_drop_while_func, + (VALUE) memo); +} + +static VALUE lazy_lazy(VALUE obj) { return obj; @@ -1524,11 +1591,15 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "grep", lazy_grep, 1); rb_define_method(rb_cLazy, "zip", lazy_zip, -1); rb_define_method(rb_cLazy, "take", lazy_take, 1); + rb_define_method(rb_cLazy, "take_while", lazy_take_while, 0); + rb_define_method(rb_cLazy, "drop", lazy_drop, 1); + rb_define_method(rb_cLazy, "drop_while", lazy_drop_while, 0); rb_define_method(rb_cLazy, "lazy", lazy_lazy, 0); rb_define_alias(rb_cLazy, "collect", "map"); rb_define_alias(rb_cLazy, "collect_concat", "flat_map"); rb_define_alias(rb_cLazy, "find_all", "select"); + rb_define_alias(rb_cLazy, "force", "to_a"); rb_eStopIteration = rb_define_class("StopIteration", rb_eIndexError); rb_define_method(rb_eStopIteration, "result", stop_result, 0); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 3adaaa32f3..6b4d49a041 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -145,10 +145,46 @@ class TestLazyEnumerator < Test::Unit::TestCase end def test_take - a = Step.new(1..3) - assert_equal(1, a.take(2).first) - assert_equal(2, a.current) - assert_equal(1, a.lazy.take(2).first) + a = Step.new(1..10) + assert_equal(1, a.take(5).first) + assert_equal(5, a.current) + assert_equal(1, a.lazy.take(5).first) + assert_equal(1, a.current) + assert_equal((1..5).to_a, a.lazy.take(5).to_a) + end + + def test_take_while + a = Step.new(1..10) + assert_equal(1, a.take_while {|i| i < 5}.first) + assert_equal(5, a.current) + assert_equal(1, a.lazy.take_while {|i| i < 5}.first) assert_equal(1, a.current) + assert_equal((1..4).to_a, a.lazy.take_while {|i| i < 5}.to_a) + end + + def test_drop + a = Step.new(1..10) + assert_equal(6, a.drop(5).first) + assert_equal(10, a.current) + assert_equal(6, a.lazy.drop(5).first) + assert_equal(6, a.current) + assert_equal((6..10).to_a, a.lazy.drop(5).to_a) + end + + def test_drop_while + a = Step.new(1..10) + assert_equal(5, a.drop_while {|i| i < 5}.first) + assert_equal(10, a.current) + assert_equal(5, a.lazy.drop_while {|i| i < 5}.first) + assert_equal(5, a.current) + assert_equal((5..10).to_a, a.lazy.drop_while {|i| i < 5}.to_a) + end + + def test_drop_and_take + assert_equal([4, 5], (1..Float::INFINITY).lazy.drop(3).take(2).to_a) + end + + def test_force + assert_equal([1, 2, 3], (1..Float::INFINITY).lazy.take(3).force) end end |