Fix bugs in Peekable and Flatten when using non-fused iterators

src/libcore/iter/adapters/flatten.rs

@@ -1,7 +1,7 @@
 use crate::fmt;
 use crate::ops::Try;
 
-use super::super::{DoubleEndedIterator, FusedIterator, Iterator};
+use super::super::{DoubleEndedIterator, Fuse, FusedIterator, Iterator};
 use super::Map;
 
 /// An iterator that maps each element to an iterator, and yields the elements
@@ -239,14 +239,17 @@ where
 /// this type.
 #[derive(Clone, Debug)]
 struct FlattenCompat<I, U> {
-    iter: I,
+    iter: Fuse<I>,
     frontiter: Option<U>,
     backiter: Option<U>,
 }
-impl<I, U> FlattenCompat<I, U> {
+impl<I, U> FlattenCompat<I, U>
+where
+    I: Iterator,
+{
     /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
     fn new(iter: I) -> FlattenCompat<I, U> {
-        FlattenCompat { iter, frontiter: None, backiter: None }
+        FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None }
     }
 }
 
@@ -261,8 +264,9 @@ where
     fn next(&mut self) -> Option<U::Item> {
         loop {
             if let Some(ref mut inner) = self.frontiter {
-                if let elt @ Some(_) = inner.next() {
-                    return elt;
+                match inner.next() {
+                    None => self.frontiter = None,
+                    elt @ Some(_) => return elt,
                 }
             }
             match self.iter.next() {
@@ -348,8 +352,9 @@ where
     fn next_back(&mut self) -> Option<U::Item> {
         loop {
             if let Some(ref mut inner) = self.backiter {
-                if let elt @ Some(_) = inner.next_back() {
-                    return elt;
+                match inner.next_back() {
+                    None => self.backiter = None,
+                    elt @ Some(_) => return elt,
                 }
             }
             match self.iter.next_back() {

src/libcore/iter/adapters/mod.rs

@@ -1468,7 +1468,11 @@ where
 {
     #[inline]
     fn next_back(&mut self) -> Option<Self::Item> {
-        self.iter.next_back().or_else(|| self.peeked.take().and_then(|x| x))
+        match self.peeked.as_mut() {
+            Some(v @ Some(_)) => self.iter.next_back().or_else(|| v.take()),
+            Some(None) => None,
+            None => self.iter.next_back(),
+        }
     }
 
     #[inline]

src/libcore/tests/iter.rs

@@ -1,3 +1,5 @@
+// ignore-tidy-filelength
+
 use core::cell::Cell;
 use core::convert::TryFrom;
 use core::iter::*;
@@ -2940,3 +2942,73 @@ fn test_partition() {
     check(xs, |&x| x < 3, 3); // small
     check(xs, |&x| x > 6, 3); // large
 }
+
+/// An iterator that panics whenever `next` or next_back` is called
+/// after `None` has already been returned. This does not violate
+/// `Iterator`'s contract. Used to test that iterator adaptors don't
+/// poll their inner iterators after exhausting them.
+struct NonFused<I> {
+    iter: I,
+    done: bool,
+}
+
+impl<I> NonFused<I> {
+    fn new(iter: I) -> Self {
+        Self { iter, done: false }
+    }
+}
+
+impl<I> Iterator for NonFused<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        assert!(!self.done, "this iterator has already returned None");
+        self.iter.next().or_else(|| {
+            self.done = true;
+            None
+        })
+    }
+}
+
+impl<I> DoubleEndedIterator for NonFused<I>
+where
+    I: DoubleEndedIterator,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        assert!(!self.done, "this iterator has already returned None");
+        self.iter.next_back().or_else(|| {
+            self.done = true;
+            None
+        })
+    }
+}
+
+#[test]
+fn test_peekable_non_fused() {
+    let mut iter = NonFused::new(empty::<i32>()).peekable();
+
+    assert_eq!(iter.peek(), None);
+    assert_eq!(iter.next_back(), None);
+}
+
+#[test]
+fn test_flatten_non_fused_outer() {
+    let mut iter = NonFused::new(once(0..2)).flatten();
+
+    assert_eq!(iter.next_back(), Some(1));
+    assert_eq!(iter.next(), Some(0));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn test_flatten_non_fused_inner() {
+    let mut iter = once(0..1).chain(once(1..3)).flat_map(NonFused::new);
+
+    assert_eq!(iter.next_back(), Some(2));
+    assert_eq!(iter.next(), Some(0));
+    assert_eq!(iter.next(), Some(1));
+    assert_eq!(iter.next(), None);
+}