Skip to content

Commit

Permalink
Fix Enumerator::Lazy#{to_enum,enum_for} where method is defined in Lazy
Browse files Browse the repository at this point in the history 
Previously, passing to_enum/enum_for a method that was defined in
Lazy itself returned wrong results:

  [1,2,3].to_enum(:map).to_a
  # => [1, 2, 3]
  [1,2,3].lazy.to_enum(:map).to_a
  # => []

I'm not sure why methods that are designed to be lazy do not work
with to_enum/enum_for.  However, one possible way to work around
this bug is to have to_enum/enum_for use the implementation found
in Enumerable/Enumerator, which is what this commit does.

While this commit works around the problem, it is a band-aid, not a
real fix.  It doesn't handle aliases of Enumerable::Lazy methods,
for instance.  A better fix would be appreciated.
  • Loading branch information
@jeremyevans
jeremyevans committed Sep 3, 2019
1 parent 4a3972c commit a7d7a0d
Show file tree
Hide file tree
Showing 2 changed files with 116 additions and 1 deletion.
69 changes: 68 additions & 1 deletion enumerator.c
View file
Expand Up @@ -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
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
48 changes: 48 additions & 0 deletions test/ruby/test_lazy_enumerator.rb
View file
Expand Up @@ -466,6 +466,54 @@ def lazy.foo(*args)
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
Expand Down

0 comments on commit a7d7a0d

Please sign in to comment.