diff --git a/enumerator.c b/enumerator.c index 092e950..1414961 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1423,7 +1423,7 @@ lazy_set_method(VALUE lazy, VALUE args, rb_enumerator_size_func *size_fn) * e.lazy -> lazy_enumerator * * Returns a lazy enumerator, whose methods map/collect, - * flat_map/collect_concat, select/find_all, reject, grep, zip, take, + * flat_map/collect_concat, select/find_all, reject, grep, grep_v, zip, take, * take_while, drop, and drop_while enumerate values only on an * as-needed basis. However, if a block is given to zip, values * are enumerated immediately. @@ -1691,6 +1691,40 @@ lazy_grep(VALUE obj, VALUE pattern) } static VALUE +lazy_grep_v_func(RB_BLOCK_CALL_FUNC_ARGLIST(val, m)) +{ + 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, i); + } + return Qnil; +} + +static VALUE +lazy_grep_v_iter(RB_BLOCK_CALL_FUNC_ARGLIST(val, m)) +{ + 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_v(VALUE obj, VALUE pattern) +{ + return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj, + rb_block_given_p() ? + lazy_grep_v_iter : lazy_grep_v_func, + pattern), + rb_ary_new3(1, pattern), 0); +} + +static VALUE call_next(VALUE obj) { return rb_funcall(obj, id_next, 0); @@ -2029,6 +2063,7 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "find_all", lazy_select, 0); rb_define_method(rb_cLazy, "reject", lazy_reject, 0); rb_define_method(rb_cLazy, "grep", lazy_grep, 1); + rb_define_method(rb_cLazy, "grep_v", lazy_grep_v, 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); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 952a773..a0c5068 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -197,6 +197,34 @@ def test_grep_multiple_values e.lazy.grep(proc {|x| x == [2, "2"]}, &:join).force) end + def test_grep_v + a = Step.new('a'..'f') + assert_equal('b', a.grep_v(/a/).first) + assert_equal('f', a.current) + assert_equal('a', a.lazy.grep_v(/c/).first) + assert_equal('a', a.current) + assert_equal(%w[b c d f], a.grep_v(proc {|x| /[aeiou]/ =~ x})) + assert_equal(%w[b c d f], a.lazy.grep_v(proc {|x| /[aeiou]/ =~ x}).to_a) + end + + def test_grep_v_with_block + a = Step.new('a'..'f') + assert_equal('B', a.grep_v(/a/) {|i| i.upcase}.first) + assert_equal('B', a.lazy.grep_v(/a/) {|i| i.upcase}.first) + end + + def test_grep_v_multiple_values + e = Enumerator.new { |yielder| + 3.times { |i| + yielder.yield(i, i.to_s) + } + } + assert_equal([[0, "0"], [1, "1"]], e.grep_v(proc {|x| x == [2, "2"]})) + assert_equal([[0, "0"], [1, "1"]], e.lazy.grep_v(proc {|x| x == [2, "2"]}).force) + assert_equal(["00", "11"], + e.lazy.grep_v(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)