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 { - iter: I, + iter: Fuse, frontiter: Option, backiter: Option, } -impl FlattenCompat { +impl FlattenCompat +where + I: Iterator, +{ /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`. fn new(iter: I) -> FlattenCompat { - FlattenCompat { iter, frontiter: None, backiter: None } + FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None } } } @@ -261,8 +264,9 @@ where fn next(&mut self) -> Option { 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 { 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.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 { + iter: I, + done: bool, +} + +impl NonFused { + fn new(iter: I) -> Self { + Self { iter, done: false } + } +} + +impl Iterator for NonFused +where + I: Iterator, +{ + type Item = I::Item; + + fn next(&mut self) -> Option { + assert!(!self.done, "this iterator has already returned None"); + self.iter.next().or_else(|| { + self.done = true; + None + }) + } +} + +impl DoubleEndedIterator for NonFused +where + I: DoubleEndedIterator, +{ + fn next_back(&mut self) -> Option { + 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::()).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); +}