diff --git a/ChangeLog b/ChangeLog index ab57cfac7d..a7544fa261 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,3 +1,9 @@ +Wed Mar 14 08:06:35 2012 Shugo Maeda + + * enumerator.c (lazy_zip): add Enumerable::Lazy#flat_map. + + * enumerator.c (lazy_lazy): just returns self. + Wed Mar 14 07:48:36 2012 Tadayoshi Funaba * ext/date/date_core.c (datetime_s_now): [ruby-core:43256]. diff --git a/enumerator.c b/enumerator.c index 5e91673899..65e750e7ba 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1388,6 +1388,63 @@ lazy_grep(VALUE obj, VALUE pattern) return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_grep_func, pattern); } +static VALUE +lazy_zip_func_i(VALUE val, VALUE arg, int argc, VALUE *argv) +{ + VALUE yielder, ary, v, result; + int i; + + yielder = argv[0]; + ary = rb_ary_new2(RARRAY_LEN(arg) + 1); + rb_ary_push(ary, argv[1]); + for (i = 0; i < RARRAY_LEN(arg); i++) { + v = rb_funcall(RARRAY_PTR(arg)[i], rb_intern("next"), 0); + rb_ary_push(ary, v); + } + result = rb_yield(ary); + rb_funcall(yielder, id_yield, 1, result); + return Qnil; +} + +static VALUE +lazy_zip_func(VALUE val, VALUE arg, int argc, VALUE *argv) +{ + VALUE yielder, ary, v; + int i; + + yielder = argv[0]; + ary = rb_ary_new2(RARRAY_LEN(arg) + 1); + rb_ary_push(ary, argv[1]); + for (i = 0; i < RARRAY_LEN(arg); i++) { + v = rb_funcall(RARRAY_PTR(arg)[i], rb_intern("next"), 0); + rb_ary_push(ary, v); + } + rb_funcall(yielder, id_yield, 1, ary); + return Qnil; +} + +static VALUE +lazy_zip(int argc, VALUE *argv, VALUE obj) +{ + VALUE ary; + int i; + + ary = rb_ary_new2(argc); + for (i = 0; i < argc; i++) { + rb_ary_push(ary, rb_funcall(argv[i], rb_intern("lazy"), 0)); + } + + return rb_block_call(rb_cLazy, id_new, 1, &obj, + rb_block_given_p() ? lazy_zip_func_i : lazy_zip_func, + ary); +} + +static VALUE +lazy_lazy(VALUE obj) +{ + return obj; +} + static VALUE stop_result(VALUE self) { @@ -1428,6 +1485,8 @@ Init_Enumerator(void) rb_define_method(rb_cLazy, "select", 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, "zip", lazy_zip, -1); + 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"); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 9b74429a7f..6d96438454 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -109,4 +109,26 @@ class TestLazyEnumerator < Test::Unit::TestCase assert_equal('c', a.lazy.grep(/c/).first) assert_equal('c', a.current) end + + def test_zip + a = Step.new(1..3) + assert_equal([1, "a"], a.zip("a".."c").first) + assert_equal(3, a.current) + assert_equal([1, "a"], a.lazy.zip("a".."c").first) + assert_equal(1, a.current) + end + + def test_zip_without_arg + a = Step.new(1..3) + assert_equal([1], a.zip.first) + assert_equal(3, a.current) + assert_equal([1], a.lazy.zip.first) + assert_equal(1, a.current) + end + + def test_zip_with_block + a = Step.new(1..3) + assert_equal(["a", 1], a.lazy.zip("a".."c") {|x, y| [y, x]}.first) + assert_equal(1, a.current) + end end