core::iter: Peekable should remember peeking a None

bluss · bluss · commit 6c2a456bf198 · 2016-11-17T18:30:00.000+01:00
Peekable must remember if a None has been seen in the `.peek()` method.
It ensures that `.peek(); .peek();` or `.peek(); .next();` only advances the
underlying iterator at most once. This does not by itself make the iterator
fused.
diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
@@ -1273,54 +1273,68 @@ unsafe impl<I> TrustedLen for Enumerate<I>
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Peekable<I: Iterator> {
     iter: I,
-    peeked: Option<I::Item>,
+    /// Remember a peeked value, even if it was None.
+    peeked: Option<Option<I::Item>>,
 }
 
+// Peekable must remember if a None has been seen in the `.peek()` method.
+// It ensures that `.peek(); .peek();` or `.peek(); .next();` only advances the
+// underlying iterator at most once. This does not by itself make the iterator
+// fused.
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<I: Iterator> Iterator for Peekable<I> {
     type Item = I::Item;
 
     #[inline]
     fn next(&mut self) -> Option<I::Item> {
-        match self.peeked {
-            Some(_) => self.peeked.take(),
+        match self.peeked.take() {
+            Some(v) => v,
             None => self.iter.next(),
         }
     }
 
     #[inline]
     #[rustc_inherit_overflow_checks]
-    fn count(self) -> usize {
-        (if self.peeked.is_some() { 1 } else { 0 }) + self.iter.count()
+    fn count(mut self) -> usize {
+        match self.peeked.take() {
+            Some(None) => 0,
+            Some(Some(_)) => 1 + self.iter.count(),
+            None => self.iter.count(),
+        }
     }
 
     #[inline]
     fn nth(&mut self, n: usize) -> Option<I::Item> {
-        match self.peeked {
-            Some(_) if n == 0 => self.peeked.take(),
-            Some(_) => {
-                self.peeked = None;
-                self.iter.nth(n-1)
-            },
-            None => self.iter.nth(n)
+        match self.peeked.take() {
+            // the .take() below is just to avoid "move into pattern guard"
+            Some(ref mut v) if n == 0 => v.take(),
+            Some(None) => None,
+            Some(Some(_)) => self.iter.nth(n - 1),
+            None => self.iter.nth(n),
         }
     }
 
     #[inline]
-    fn last(self) -> Option<I::Item> {
-        self.iter.last().or(self.peeked)
+    fn last(mut self) -> Option<I::Item> {
+        let peek_opt = match self.peeked.take() {
+            Some(None) => return None,
+            Some(v) => v,
+            None => None,
+        };
+        self.iter.last().or(peek_opt)
     }
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
+        let peek_len = match self.peeked {
+            Some(None) => return (0, Some(0)),
+            Some(Some(_)) => 1,
+            None => 0,
+        };
         let (lo, hi) = self.iter.size_hint();
-        if self.peeked.is_some() {
-            let lo = lo.saturating_add(1);
-            let hi = hi.and_then(|x| x.checked_add(1));
-            (lo, hi)
-        } else {
-            (lo, hi)
-        }
+        let lo = lo.saturating_add(peek_len);
+        let hi = hi.and_then(|x| x.checked_add(peek_len));
+        (lo, hi)
     }
 }
 
@@ -1372,9 +1386,13 @@ impl<I: Iterator> Peekable<I> {
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn peek(&mut self) -> Option<&I::Item> {
         if self.peeked.is_none() {
-            self.peeked = self.iter.next();
+            self.peeked = Some(self.iter.next());
+        }
+        match self.peeked {
+            Some(Some(ref value)) => Some(value),
+            Some(None) => None,
+            _ => unreachable!(),
         }
-        self.peeked.as_ref()
     }
 }
 
diff --git a/src/libcoretest/iter.rs b/src/libcoretest/iter.rs
@@ -274,6 +274,74 @@ fn test_iterator_peekable_last() {
     let mut it = ys.iter().peekable();
     assert_eq!(it.peek(), Some(&&0));
     assert_eq!(it.last(), Some(&0));
+
+    let mut it = ys.iter().peekable();
+    assert_eq!(it.next(), Some(&0));
+    assert_eq!(it.peek(), None);
+    assert_eq!(it.last(), None);
+}
+
+/// This is an iterator that follows the Iterator contract,
+/// but it is not fused. After having returned None once, it will start
+/// producing elements if .next() is called again.
+pub struct CycleIter<'a, T: 'a> {
+    index: usize,
+    data: &'a [T],
+}
+
+pub fn cycle<T>(data: &[T]) -> CycleIter<T> {
+    CycleIter {
+        index: 0,
+        data: data,
+    }
+}
+
+impl<'a, T> Iterator for CycleIter<'a, T> {
+    type Item = &'a T;
+    fn next(&mut self) -> Option<Self::Item> {
+        let elt = self.data.get(self.index);
+        self.index += 1;
+        self.index %= 1 + self.data.len();
+        elt
+    }
+}
+
+#[test]
+fn test_iterator_peekable_remember_peek_none_1() {
+    // Check that the loop using .peek() terminates
+    let data = [1, 2, 3];
+    let mut iter = cycle(&data).peekable();
+
+    let mut n = 0;
+    while let Some(_) = iter.next() {
+        let is_the_last = iter.peek().is_none();
+        assert_eq!(is_the_last, n == data.len() - 1);
+        n += 1;
+        if n > data.len() { break; }
+    }
+    assert_eq!(n, data.len());
+}
+
+#[test]
+fn test_iterator_peekable_remember_peek_none_2() {
+    let data = [0];
+    let mut iter = cycle(&data).peekable();
+    iter.next();
+    assert_eq!(iter.peek(), None);
+    assert_eq!(iter.last(), None);
+}
+
+#[test]
+fn test_iterator_peekable_remember_peek_none_3() {
+    let data = [0];
+    let mut iter = cycle(&data).peekable();
+    iter.peek();
+    assert_eq!(iter.nth(0), Some(&0));
+
+    let mut iter = cycle(&data).peekable();
+    iter.next();
+    assert_eq!(iter.peek(), None);
+    assert_eq!(iter.nth(0), None);
 }
 
 #[test]
