Iterator::Flatten#rewind returns iterator itself

crystal-lang · Dec 23, 2016 · 352d5f0 · 352d5f0
1 parent 0dfd469
commit 352d5f0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/spec/std/iterator_spec.cr b/spec/std/iterator_spec.cr
@@ -632,5 +632,12 @@ describe Iterator do
       iter.next.should eq(7)
       iter.next.should be_a(Iterator::Stop)
     end
+
+    it "return iterator itself by rewind" do
+      iter = [1, [2, 3], 4].each.flatten
+
+      iter.to_a.should eq([1, 2, 3, 4])
+      iter.rewind.to_a.should eq([1, 2, 3, 4])
+    end
   end
 end
diff --git a/src/iterator.cr b/src/iterator.cr
@@ -448,6 +448,7 @@ module Iterator(T)
       @generators << @iterator
       @stopped.each &.rewind
       @stopped.clear
+      self
     end
 
     def self.element_type(element)