diff options
| -rw-r--r-- | enumerator.c | 69 | ||||
| -rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 48 |
2 files changed, 116 insertions, 1 deletions
diff --git a/enumerator.c b/enumerator.c index d4022aefc4..62534ee241 100644 --- a/enumerator.c +++ b/enumerator.c @@ -113,6 +113,8 @@ static ID id_next, id_result, id_receiver, id_arguments, id_memo, id_method, id_ static ID id_begin, id_end, id_step, id_exclude_end; static VALUE sym_each, sym_cycle, sym_yield; +static VALUE lazy_use_super_method; + #define id_call idCall #define id_each idEach #define id_eqq idEqq @@ -1875,12 +1877,15 @@ lazy_to_enum_i(VALUE obj, VALUE meth, int argc, const VALUE *argv, rb_enumerator static VALUE lazy_to_enum(int argc, VALUE *argv, VALUE self) { - VALUE lazy, meth = sym_each; + VALUE lazy, meth = sym_each, super_meth; if (argc > 0) { --argc; meth = *argv++; } + if (RTEST((super_meth = rb_hash_aref(lazy_use_super_method, meth)))) { + meth = super_meth; + } lazy = lazy_to_enum_i(self, meth, argc, argv, 0); if (rb_block_given_p()) { enumerator_ptr(lazy)->size = rb_block_proc(); @@ -3665,6 +3670,8 @@ lazy_with_index(int argc, VALUE *argv, VALUE obj) void InitVM_Enumerator(void) { + ID id_private = rb_intern("private"); + rb_define_method(rb_mKernel, "to_enum", obj_to_enum, -1); rb_define_method(rb_mKernel, "enum_for", obj_to_enum, -1); @@ -3693,6 +3700,44 @@ InitVM_Enumerator(void) /* Lazy */ rb_cLazy = rb_define_class_under(rb_cEnumerator, "Lazy", rb_cEnumerator); rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0); + + rb_define_alias(rb_cLazy, "_enumerable_map", "map"); + rb_define_alias(rb_cLazy, "_enumerable_collect", "collect"); + rb_define_alias(rb_cLazy, "_enumerable_flat_map", "flat_map"); + rb_define_alias(rb_cLazy, "_enumerable_collect_concat", "collect_concat"); + rb_define_alias(rb_cLazy, "_enumerable_select", "select"); + rb_define_alias(rb_cLazy, "_enumerable_find_all", "find_all"); + rb_define_alias(rb_cLazy, "_enumerable_filter", "filter"); + rb_define_alias(rb_cLazy, "_enumerable_filter_map", "filter_map"); + rb_define_alias(rb_cLazy, "_enumerable_reject", "reject"); + rb_define_alias(rb_cLazy, "_enumerable_grep", "grep"); + rb_define_alias(rb_cLazy, "_enumerable_grep_v", "grep_v"); + rb_define_alias(rb_cLazy, "_enumerable_zip", "zip"); + rb_define_alias(rb_cLazy, "_enumerable_take", "take"); + rb_define_alias(rb_cLazy, "_enumerable_take_while", "take_while"); + rb_define_alias(rb_cLazy, "_enumerable_drop", "drop"); + rb_define_alias(rb_cLazy, "_enumerable_drop_while", "drop_while"); + rb_define_alias(rb_cLazy, "_enumerable_uniq", "uniq"); + rb_define_private_method(rb_cLazy, "_enumerable_with_index", enumerator_with_index, -1); + + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_map"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_collect"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_flat_map"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_collect_concat"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_select"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_find_all"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_filter"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_filter_map"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_reject"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_grep"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_grep_v"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_zip"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_take"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_take_while"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_drop"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_drop_while"))); + rb_funcall(rb_cLazy, id_private, 1, ID2SYM(rb_intern("_enumerable_uniq"))); + rb_define_method(rb_cLazy, "initialize", lazy_initialize, -1); rb_define_method(rb_cLazy, "to_enum", lazy_to_enum, -1); rb_define_method(rb_cLazy, "enum_for", lazy_to_enum, -1); @@ -3721,6 +3766,28 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "uniq", lazy_uniq, 0); rb_define_method(rb_cLazy, "with_index", lazy_with_index, -1); + lazy_use_super_method = rb_hash_new(); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("map")), ID2SYM(rb_intern("_enumerable_map"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("collect")), ID2SYM(rb_intern("_enumerable_collect"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("flat_map")), ID2SYM(rb_intern("_enumerable_flat_map"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("collect_concat")), ID2SYM(rb_intern("_enumerable_collect_concat"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("select")), ID2SYM(rb_intern("_enumerable_select"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("find_all")), ID2SYM(rb_intern("_enumerable_find_all"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("filter")), ID2SYM(rb_intern("_enumerable_filter"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("filter_map")), ID2SYM(rb_intern("_enumerable_filter_map"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("reject")), ID2SYM(rb_intern("_enumerable_reject"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("grep")), ID2SYM(rb_intern("_enumerable_grep"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("grep_v")), ID2SYM(rb_intern("_enumerable_grep_v"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("zip")), ID2SYM(rb_intern("_enumerable_zip"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("take")), ID2SYM(rb_intern("_enumerable_take"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("take_while")), ID2SYM(rb_intern("_enumerable_take_while"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("drop")), ID2SYM(rb_intern("_enumerable_drop"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("drop_while")), ID2SYM(rb_intern("_enumerable_drop_while"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("uniq")), ID2SYM(rb_intern("_enumerable_uniq"))); + rb_hash_aset(lazy_use_super_method, ID2SYM(rb_intern("with_index")), ID2SYM(rb_intern("_enumerable_with_index"))); + rb_obj_freeze(lazy_use_super_method); + rb_gc_register_mark_object(lazy_use_super_method); + #if 0 /* for RDoc */ rb_define_method(rb_cLazy, "to_a", lazy_to_a, 0); rb_define_method(rb_cLazy, "chunk", lazy_chunk, 0); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 80977a0bbd..72fe1a4a24 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -466,6 +466,54 @@ EOS assert_equal [1, 2, 3], lazy.to_enum.to_a end + def test_lazy_to_enum_lazy_methods + a = [1, 2, 3].to_enum + pr = proc{|x| [x, x * 2]} + selector = proc{|x| x*2 if x % 2 == 0} + + [ + [:with_index, nil], + [:with_index, 10, nil], + [:with_index, 10, pr], + [:map, nil], + [:map, pr], + [:collect, nil], + [:flat_map, nil], + [:flat_map, pr], + [:collect_concat, nil], + [:select, nil], + [:select, selector], + [:find_all, nil], + [:filter, nil], + [:filter_map, selector], + [:filter_map, nil], + [:reject, selector], + [:grep, selector, nil], + [:grep, selector, pr], + [:grep_v, selector, nil], + [:grep_v, selector, pr], + [:zip, a, nil], + [:take, 3, nil], + [:take_while, nil], + [:take_while, selector], + [:drop, 1, nil], + [:drop_while, nil], + [:drop_while, selector], + [:uniq, nil], + [:uniq, proc{|x| x.odd?}], + ].each do |args| + block = args.pop + assert_equal [1, 2, 3].to_enum.to_enum(*args).first(2).to_a, [1, 2, 3].to_enum.lazy.to_enum(*args).first(2).to_a + assert_equal (0..50).to_enum.to_enum(*args).first(2).to_a, (0..50000).to_enum.lazy.to_enum(*args).first(2).to_a + if block + assert_equal [1, 2, 3, 4].to_enum.to_enum(*args).map(&block).first(2).to_a, [1, 2, 3, 4].to_enum.lazy.to_enum(*args).map(&block).first(2).to_a + unless args.first == :take_while || args.first == :drop_while + assert_equal (0..50).to_enum.to_enum(*args).map(&block).first(2).to_a, (0..50000).to_enum.lazy.to_enum(*args).map(&block).first(2).to_a + end + end + end + end + def test_size lazy = [1, 2, 3].lazy assert_equal 3, lazy.size |