Auto merge of #68915 - timvermeulen:non_fused_iter, r=Amanieu

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

I fixed a couple of bugs with regard to the `Peekable` and `Flatten`/`FlatMap` iterators when the underlying iterator isn't fused. For testing, I also added a `NonFused` iterator wrapper that panics when `next` or `next_back` is called on an iterator that has returned `None` before, which will hopefully make it easier to spot these mistakes in the future.

### Peekable

`Peekable::next_back` was implemented as
```rust
self.iter.next_back().or_else(|| self.peeked.take().and_then(|x| x))
```
which is incorrect because when the `peeked` field is `Some(None)`, then `None` has already been returned from the inner iterator and what it returns from `next_back` can no longer be relied upon. `test_peekable_non_fused` tests this.

### Flatten

When a `FlattenCompat` instance only has a `backiter` remaining (i.e. `self.frontiter` is `None` and `self.iter` is empty), then `next` will call `self.iter.next()` every time, so the `iter` field needs to be fused. I fixed it by giving it the type `Fuse<I>` instead of `I`, I think this is the only way to fix it. `test_flatten_non_fused_outer` tests this.

Furthermore, previously `FlattenCompat::next` did not set `self.frontiter` to `None` after it returned `None`, which is incorrect when the inner iterator type isn't fused. I just delegated it to `try_fold` because that already handles it correctly. `test_flatten_non_fused_inner` tests this.

r? @scottmcm

3 files changed, 90 insertions, 9 deletions

diff --git a/src/libcore/iter/adapters/flatten.rs b/src/libcore/iter/adapters/flatten.rs
index 0a7a9f26f89..4202e52448d 100644
--- a/src/libcore/iter/adapters/flatten.rs
+++ b/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() {
diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs
index 26132e36c97..3c0ddcb2bc8 100644
--- a/src/libcore/iter/adapters/mod.rs
+++ b/src/libcore/iter/adapters/mod.rs
@@ -1480,7 +1480,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]
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 5b41ef35065..98e3eeb982b 100644
--- a/src/libcore/tests/iter.rs
+++ b/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);
+}