diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs index 29f54bd10fba1..1ef7380c8e258 100644 --- a/src/libstd/iterator.rs +++ b/src/libstd/iterator.rs @@ -221,6 +221,28 @@ pub trait IteratorUtil { /// ~~~ fn enumerate(self) -> Enumerate; + + /// Creates an iterator that has a `.peek()` method + /// that returns a optional reference to the next element, + /// and a `.peek_clone()` method that returns an optional copy. + /// + /// # Example + /// + /// ~~~ {.rust} + /// let a = [100, 200, 300]; + /// let mut it = a.consume_iter().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.peek_clone().unwrap(), 300); + /// assert_eq!(it.next().unwrap(), 300); + /// assert!(it.peek().is_none()); + /// assert!(it.next().is_none()); + /// ~~~ + fn peekable(self) -> Peekable; + /// Creates an iterator which invokes the predicate on elements until it /// returns false. Once the predicate returns false, all further elements are /// yielded. @@ -520,6 +542,11 @@ impl> IteratorUtil for T { Enumerate{iter: self, count: 0} } + #[inline] + fn peekable(self) -> Peekable { + Peekable{iter: self, peeked: None} + } + #[inline] fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhile<'r, A, T> { SkipWhile{iter: self, flag: false, predicate: predicate} @@ -1154,6 +1181,56 @@ impl> RandomAccessIterator<(uint, A)> for Enumerat } } +/// An iterator with a `peek()` that returns an optional reference to the next element. +pub struct Peekable { + priv iter: T, + priv peeked: Option, +} + +impl> Iterator for Peekable { + #[inline] + fn next(&mut self) -> Option { + if self.peeked.is_some() { self.peeked.take() } + else { self.iter.next() } + } +} + +impl<'self, A, T: Iterator> Peekable { + /// 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, + } + }, + } + } +} + +impl> Peekable { + /// Return a copy of the next element of the iterator with out advancing it, + /// or None if the iterator is exhausted. + #[inline] + pub fn peek_clone(&mut self) -> Option { + match self.peeked { + Some(ref value) => Some(value.clone()), + None => { + self.peeked = self.iter.next(); + match self.peeked { + Some(ref value) => Some(value.clone()), + None => None, + } + }, + } + } +} + /// An iterator which rejects elements while `predicate` is true pub struct SkipWhile<'self, A, T> { priv iter: T, @@ -1645,6 +1722,25 @@ mod tests { } } + #[test] + fn test_iterator_peekable() { + let xs = ~[0u, 1, 2, 3, 4, 5]; + let mut it = xs.consume_iter().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.peek_clone().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];