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

  [bug #7696]

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

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@38920 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

1 files changed, 17 insertions, 0 deletions

diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index acd4843afb..35e92c9050 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -243,6 +243,23 @@ class TestLazyEnumerator < Test::Unit::TestCase
     assert_equal((1..5).to_a, take5.force, bug6428)
   end
 
+  def test_take_nested
+    bug7696 = '[ruby-core:51470]'
+    a = Step.new(1..10)
+    take5 = a.lazy.take(5)
+    assert_equal([*(1..5)]*5, take5.flat_map{take5}.force, bug7696)
+  end
+
+  def test_take_rewound
+    bug7696 = '[ruby-core:51470]'
+    e=(1..42).lazy.take(2)
+    assert_equal 1, e.next
+    assert_equal 2, e.next
+    e.rewind
+    assert_equal 1, e.next
+    assert_equal 2, e.next
+  end
+
   def test_take_while
     a = Step.new(1..10)
     assert_equal(1, a.take_while {|i| i < 5}.first)