diff options
| author | shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-03-09 05:34:41 +0000 |
|---|---|---|
| committer | shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-03-09 05:34:41 +0000 |
| commit | a21d0f72c280a3effe15554279ae006918cd97ce (patch) | |
| tree | 4656f544b87a81044c5b0192be5d82bb9289bf2f | |
| parent | 5452b18f9905ea404cd0bbfff5536f031bb83450 (diff) | |
| download | ruby-a21d0f72c280a3effe15554279ae006918cd97ce.tar.gz | |
* enumerator.c (lazy_flat_map): add Enumerable::Lazy#flat_map.
git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@34956 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
| -rw-r--r-- | ChangeLog | 4 | ||||
| -rw-r--r-- | enumerator.c | 28 | ||||
| -rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 8 |
3 files changed, 40 insertions, 0 deletions
@@ -1,3 +1,7 @@ +Fri Mar 9 14:29:32 2012 Shugo Maeda <shugo@ruby-lang.org> + + * enumerator.c (lazy_flat_map): add Enumerable::Lazy#flat_map. + Fri Mar 9 06:29:22 2012 Aaron Patterson <aaron@tenderlovemaking.com> * ext/psych/lib/psych.rb (load, parse): stop parsing or loading after diff --git a/enumerator.c b/enumerator.c index c1a553a705..9e3da15dd3 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1272,6 +1272,32 @@ lazy_map(VALUE obj) return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_map_func, 0); } +static VALUE +lazy_flat_map_func(VALUE val, VALUE m, int argc, VALUE *argv) +{ + VALUE result = rb_yield_values2(argc - 1, &argv[1]); + VALUE ary = rb_check_array_type(result); + if (NIL_P(ary)) { + return rb_funcall(argv[0], id_yield, 1, result); + } + else { + int i; + for (i = 0; i < RARRAY_LEN(ary); i++) { + rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(ary)[i]); + } + return Qnil; + } +} + +static VALUE +lazy_flat_map(VALUE obj) +{ + if (!rb_block_given_p()) { + rb_raise(rb_eArgError, "tried to call lazy flat_map without a block"); + } + + return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_flat_map_func, 0); +} static VALUE lazy_select_func(VALUE val, VALUE m, int argc, VALUE *argv) @@ -1377,11 +1403,13 @@ Init_Enumerator(void) rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0); rb_define_method(rb_cLazy, "initialize", lazy_initialize, 1); rb_define_method(rb_cLazy, "map", lazy_map, 0); + rb_define_method(rb_cLazy, "flat_map", lazy_flat_map, 0); 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_alias(rb_cLazy, "collect", "map"); + rb_define_alias(rb_cLazy, "collect_concat", "flat_map"); rb_define_alias(rb_cLazy, "find_all", "select"); rb_eStopIteration = rb_define_class("StopIteration", rb_eIndexError); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index eeede6600c..9b74429a7f 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -80,6 +80,14 @@ class TestLazyEnumerator < Test::Unit::TestCase assert_equal(1, a.current) end + def test_flat_map + a = Step.new(1..3) + assert_equal(2, a.flat_map {|x| [x * 2]}.first) + assert_equal(3, a.current) + assert_equal(2, a.lazy.flat_map {|x| [x * 2]}.first) + assert_equal(1, a.current) + end + def test_reject a = Step.new(1..6) assert_equal(4, a.reject {|x| x < 4}.first) |