diff --git a/enumerator.c b/enumerator.c index 2563de5..3d26b0b 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1533,6 +1533,7 @@ lazy_take_func(VALUE val, VALUE args, int argc, VALUE *argv) rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); if (--memo->u3.cnt == 0) { + memo->u3.cnt = memo->u2.argc; return Qundef; } else { @@ -1557,7 +1558,7 @@ lazy_take(VALUE obj, VALUE n) argv[2] = INT2NUM(0); argc = 3; } - memo = NEW_MEMO(0, 0, len); + memo = NEW_MEMO(0, len, len); return lazy_set_method(rb_block_call(rb_cLazy, id_new, argc, argv, lazy_take_func, (VALUE) memo), rb_ary_new3(1, n)); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 6920ace..064cac4 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -229,6 +229,13 @@ class TestLazyEnumerator < Test::Unit::TestCase assert_equal(nil, a.current) end + def test_take_recycle + a = Step.new(1..10) + take5 = a.lazy.take(5) + assert_equal((1..5).to_a, take5.force) + assert_equal((1..5).to_a, take5.force) + end + def test_take_while a = Step.new(1..10) assert_equal(1, a.take_while {|i| i < 5}.first)