diff --git a/src/libcore/iter/adapters/flatten.rs b/src/libcore/iter/adapters/flatten.rs index 0a7a9f26f8912..4202e52448dcf 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 7d10ef3d28219..67a68cfea083b 100644 --- a/src/libcore/iter/adapters/mod.rs +++ b/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] diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index 5b41ef350657f..98e3eeb982bde 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); +}