Add a "peekable" iterator adaptor, with a peek() method that returns the next element.

1 files changed, 72 insertions, 0 deletions

diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index 03d8ae67998..e2d157c2a59 100644
--- a/src/libstd/iterator.rs
+++ b/src/libstd/iterator.rs
@@ -156,6 +156,28 @@ pub trait Iterator<A> {
         Enumerate{iter: self, count: 0}
     }
 
+
+    /// Creates an iterator that has a `.peek()` method
+    /// that returns a optional reference to the next element.
+    ///
+    /// # Example
+    ///
+    /// ~~~ {.rust}
+    /// let a = [100, 200, 300];
+    /// let mut it = xs.iter().map(|&x|x).peekable();
+    /// assert_eq!(it.peek().unwrap(), &100);
+    /// assert_eq!(it.next().unwrap(), 100);
+    /// assert_eq!(it.next().unwrap(), 200);
+    /// assert_eq!(it.peek().unwrap(), &300);
+    /// assert_eq!(it.peek().unwrap(), &300);
+    /// assert_eq!(it.next().unwrap(), 300);
+    /// assert!(it.peek().is_none());
+    /// assert!(it.next().is_none());
+    /// ~~~
+    fn peekable(self) -> Peekable<A, Self> {
+        Peekable{iter: self, peeked: None}
+    }
+
     /// Creates an iterator which invokes the predicate on elements until it
     /// returns false. Once the predicate returns false, all further elements are
     /// yielded.
@@ -1059,6 +1081,38 @@ impl<A, T: RandomAccessIterator<A>> RandomAccessIterator<(uint, A)> for Enumerat
     }
 }
 
+/// An iterator with a `peek()` that returns an optional reference to the next element.
+pub struct Peekable<A, T> {
+    priv iter: T,
+    priv peeked: Option<A>,
+}
+
+impl<A, T: Iterator<A>> Iterator<A> for Peekable<A, T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        if self.peeked.is_some() { self.peeked.take() }
+        else { self.iter.next() }
+    }
+}
+
+impl<'self, A, T: Iterator<A>> Peekable<A, T> {
+    /// Return a reference to the next element of the iterator with out advancing it,
+    /// or None if the iterator is exhausted.
+    #[inline]
+    pub fn peek(&'self mut self) -> Option<&'self A> {
+        match self.peeked {
+            Some(ref value) => Some(value),
+            None => {
+                self.peeked = self.iter.next();
+                match self.peeked {
+                    Some(ref value) => Some(value),
+                    None => None,
+                }
+            },
+        }
+    }
+}
+
 /// An iterator which rejects elements while `predicate` is true
 pub struct SkipWhile<'self, A, T> {
     priv iter: T,
@@ -1569,6 +1623,24 @@ mod tests {
     }
 
     #[test]
+    fn test_iterator_peekable() {
+        let xs = ~[0u, 1, 2, 3, 4, 5];
+        let mut it = xs.iter().map(|&x|x).peekable();
+        assert_eq!(it.peek().unwrap(), &0);
+        assert_eq!(it.next().unwrap(), 0);
+        assert_eq!(it.next().unwrap(), 1);
+        assert_eq!(it.next().unwrap(), 2);
+        assert_eq!(it.peek().unwrap(), &3);
+        assert_eq!(it.peek().unwrap(), &3);
+        assert_eq!(it.next().unwrap(), 3);
+        assert_eq!(it.next().unwrap(), 4);
+        assert_eq!(it.peek().unwrap(), &5);
+        assert_eq!(it.next().unwrap(), 5);
+        assert!(it.peek().is_none());
+        assert!(it.next().is_none());
+    }
+
+    #[test]
     fn test_iterator_take_while() {
         let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
         let ys = [0u, 1, 2, 3, 5, 13];