* enumerator.c: Fix state handling for Lazy#drop

[bug #7696] [bug #7691]

* test/ruby/test_lazy_enumerator.rb: test for above

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@38922 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
This commit is contained in:
marcandre 2013-01-24 06:23:42 +00:00
parent bcbeb5d008
commit 9d94a1a5a1
2 changed files with 15 additions and 7 deletions

View File

@ -1725,13 +1725,16 @@ lazy_drop_size(VALUE lazy)
static VALUE static VALUE
lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv) lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv)
{ {
NODE *memo = RNODE(args); long remain;
VALUE memo = rb_ivar_get(argv[0], id_memo);
if (memo->u3.cnt == 0) { if (NIL_P(memo)) {
memo = args;
}
if ((remain = NUM2LONG(memo)) == 0) {
rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); rb_funcall2(argv[0], id_yield, argc - 1, argv + 1);
} }
else { else {
memo->u3.cnt--; rb_ivar_set(argv[0], id_memo, LONG2NUM(--remain));
} }
return Qnil; return Qnil;
} }
@ -1739,15 +1742,13 @@ lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv)
static VALUE static VALUE
lazy_drop(VALUE obj, VALUE n) lazy_drop(VALUE obj, VALUE n)
{ {
NODE *memo;
long len = NUM2LONG(n); long len = NUM2LONG(n);
if (len < 0) { if (len < 0) {
rb_raise(rb_eArgError, "attempt to drop negative size"); rb_raise(rb_eArgError, "attempt to drop negative size");
} }
memo = NEW_MEMO(0, 0, len);
return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj, return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
lazy_drop_func, (VALUE) memo), lazy_drop_func, n),
rb_ary_new3(1, n), lazy_drop_size); rb_ary_new3(1, n), lazy_drop_size);
} }

View File

@ -257,6 +257,13 @@ class TestLazyEnumerator < Test::Unit::TestCase
assert_equal([*(6..10)]*5, drop5.flat_map{drop5}.force, bug7696) assert_equal([*(6..10)]*5, drop5.flat_map{drop5}.force, bug7696)
end end
def test_drop_nested
bug7696 = '[ruby-core:51470]'
a = Step.new(1..10)
drop5 = a.lazy.drop(5)
assert_equal([*(6..10)]*5, drop5.flat_map{drop5}.force, bug7696)
end
def test_take_rewound def test_take_rewound
bug7696 = '[ruby-core:51470]' bug7696 = '[ruby-core:51470]'
e=(1..42).lazy.take(2) e=(1..42).lazy.take(2)