auto merge of #8428 : blake2-ppc/rust/peekable-iterators, r=thestinger

Peekable changes by @SimonSapin and original PR is #8396

2 files changed, 138 insertions, 107 deletions

diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs
index 424492a3cfe..f5a61692509 100644
--- a/src/libextra/treemap.rs
+++ b/src/libextra/treemap.rs
@@ -14,7 +14,8 @@
 
 
 use std::util::{swap, replace};
-use std::iterator::{FromIterator, Extendable};
+use std::iterator::{FromIterator, Extendable, Peekable};
+use std::cmp::Ordering;
 
 // This is implemented as an AA tree, which is a simplified variation of
 // a red-black tree where red (horizontal) nodes can only be added
@@ -529,24 +530,24 @@ impl<T: TotalOrd> TreeSet<T> {
 
     /// Visit the values (in-order) representing the difference
     pub fn difference<'a>(&'a self, other: &'a TreeSet<T>) -> Difference<'a, T> {
-        Difference{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
+        Difference{a: self.iter().peekable(), b: other.iter().peekable()}
     }
 
     /// Visit the values (in-order) representing the symmetric difference
     pub fn symmetric_difference<'a>(&'a self, other: &'a TreeSet<T>)
         -> SymDifference<'a, T> {
-        SymDifference{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
+        SymDifference{a: self.iter().peekable(), b: other.iter().peekable()}
     }
 
     /// Visit the values (in-order) representing the intersection
     pub fn intersection<'a>(&'a self, other: &'a TreeSet<T>)
         -> Intersection<'a, T> {
-        Intersection{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
+        Intersection{a: self.iter().peekable(), b: other.iter().peekable()}
     }
 
     /// Visit the values (in-order) representing the union
     pub fn union<'a>(&'a self, other: &'a TreeSet<T>) -> Union<'a, T> {
-        Union{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
+        Union{a: self.iter().peekable(), b: other.iter().peekable()}
     }
 }
 
@@ -560,61 +561,47 @@ pub struct TreeSetRevIterator<'self, T> {
     priv iter: TreeMapRevIterator<'self, T, ()>
 }
 
-// Encapsulate an iterator and hold its latest value until stepped forward
-struct Focus<A, T> {
-    priv iter: T,
-    priv focus: Option<A>,
-}
-
-impl<A, T: Iterator<A>> Focus<A, T> {
-    fn new(mut it: T) -> Focus<A, T> {
-        Focus{focus: it.next(), iter: it}
-    }
-    fn step(&mut self) {
-        self.focus = self.iter.next()
-    }
-}
-
 /// Lazy iterator producing elements in the set difference (in-order)
 pub struct Difference<'self, T> {
-    priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
-    priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
+    priv a: Peekable<&'self T, TreeSetIterator<'self, T>>,
+    priv b: Peekable<&'self T, TreeSetIterator<'self, T>>,
 }
 
 /// Lazy iterator producing elements in the set symmetric difference (in-order)
 pub struct SymDifference<'self, T> {
-    priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
-    priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
+    priv a: Peekable<&'self T, TreeSetIterator<'self, T>>,
+    priv b: Peekable<&'self T, TreeSetIterator<'self, T>>,
 }
 
 /// Lazy iterator producing elements in the set intersection (in-order)
 pub struct Intersection<'self, T> {
-    priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
-    priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
+    priv a: Peekable<&'self T, TreeSetIterator<'self, T>>,
+    priv b: Peekable<&'self T, TreeSetIterator<'self, T>>,
 }
 
 /// Lazy iterator producing elements in the set intersection (in-order)
 pub struct Union<'self, T> {
-    priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
-    priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
+    priv a: Peekable<&'self T, TreeSetIterator<'self, T>>,
+    priv b: Peekable<&'self T, TreeSetIterator<'self, T>>,
+}
+
+/// Compare `x` and `y`, but return `short` if x is None and `long` if y is None
+fn cmp_opt<T: TotalOrd>(x: Option<&T>, y: Option<&T>,
+                        short: Ordering, long: Ordering) -> Ordering {
+    match (x, y) {
+        (None    , _       ) => short,
+        (_       , None    ) => long,
+        (Some(x1), Some(y1)) => x1.cmp(y1),
+    }
 }
 
 impl<'self, T: TotalOrd> Iterator<&'self T> for Difference<'self, T> {
     fn next(&mut self) -> Option<&'self T> {
         loop {
-            match (self.a.focus, self.b.focus) {
-                (None    , _       ) => return None,
-                (ret     , None    ) => { self.a.step(); return ret },
-                (Some(a1), Some(b1)) => {
-                    let cmp = a1.cmp(b1);
-                    if cmp == Less {
-                        self.a.step();
-                        return Some(a1);
-                    } else {
-                        if cmp == Equal { self.a.step() }
-                        self.b.step();
-                    }
-                }
+            match cmp_opt(self.a.peek(), self.b.peek(), Less, Less) {
+                Less    => return self.a.next(),
+                Equal   => { self.a.next(); self.b.next(); }
+                Greater => { self.b.next(); }
             }
         }
     }
@@ -623,23 +610,10 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for Difference<'self, T> {
 impl<'self, T: TotalOrd> Iterator<&'self T> for SymDifference<'self, T> {
     fn next(&mut self) -> Option<&'self T> {
         loop {
-            match (self.a.focus, self.b.focus) {
-                (ret     , None    ) => { self.a.step(); return ret },
-                (None    , ret     ) => { self.b.step(); return ret },
-                (Some(a1), Some(b1)) => {
-                    let cmp = a1.cmp(b1);
-                    if cmp == Less {
-                        self.a.step();
-                        return Some(a1);
-                    } else {
-                        self.b.step();
-                        if cmp == Greater {
-                            return Some(b1);
-                        } else {
-                            self.a.step();
-                        }
-                    }
-                }
+            match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) {
+                Less    => return self.a.next(),
+                Equal   => { self.a.next(); self.b.next(); }
+                Greater => return self.b.next(),
             }
         }
     }
@@ -648,20 +622,16 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for SymDifference<'self, T> {
 impl<'self, T: TotalOrd> Iterator<&'self T> for Intersection<'self, T> {
     fn next(&mut self) -> Option<&'self T> {
         loop {
-            match (self.a.focus, self.b.focus) {
-                (None    , _       ) => return None,
-                (_       , None    ) => return None,
-                (Some(a1), Some(b1)) => {
-                    let cmp = a1.cmp(b1);
-                    if cmp == Less {
-                        self.a.step();
-                    } else {
-                        self.b.step();
-                        if cmp == Equal {
-                            return Some(a1);
-                        }
-                    }
-                },
+            let o_cmp = match (self.a.peek(), self.b.peek()) {
+                (None    , _       ) => None,
+                (_       , None    ) => None,
+                (Some(a1), Some(b1)) => Some(a1.cmp(b1)),
+            };
+            match o_cmp {
+                None          => return None,
+                Some(Less)    => { self.a.next(); }
+                Some(Equal)   => { self.b.next(); return self.a.next() }
+                Some(Greater) => { self.b.next(); }
             }
         }
     }
@@ -670,22 +640,10 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for Intersection<'self, T> {
 impl<'self, T: TotalOrd> Iterator<&'self T> for Union<'self, T> {
     fn next(&mut self) -> Option<&'self T> {
         loop {
-            match (self.a.focus, self.b.focus) {
-                (ret     , None) => { self.a.step(); return ret },
-                (None    , ret ) => { self.b.step(); return ret },
-                (Some(a1), Some(b1)) => {
-                    let cmp = a1.cmp(b1);
-                    if cmp == Greater {
-                        self.b.step();
-                        return Some(b1);
-                    } else {
-                        self.a.step();
-                        if cmp == Equal {
-                            self.b.step();
-                        }
-                        return Some(a1);
-                    }
-                }
+            match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) {
+                Less    => return self.a.next(),
+                Equal   => { self.b.next(); return self.a.next() }
+                Greater => return self.b.next(),
             }
         }
     }
diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index a7a1c0bede8..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.
@@ -286,15 +308,15 @@ pub trait Iterator<A> {
     ///let xs = [1u, 4, 2, 3, 8, 9, 6];
     ///let sum = xs.iter()
     ///            .map(|&x| x)
-    ///            .peek(|&x| debug!("filtering %u", x))
+    ///            .inspect(|&x| debug!("filtering %u", x))
     ///            .filter(|&x| x % 2 == 0)
-    ///            .peek(|&x| debug!("%u made it through", x))
+    ///            .inspect(|&x| debug!("%u made it through", x))
     ///            .sum();
     ///println(sum.to_str());
     /// ~~~
     #[inline]
-    fn peek<'r>(self, f: &'r fn(&A)) -> Peek<'r, A, Self> {
-        Peek{iter: self, f: f}
+    fn inspect<'r>(self, f: &'r fn(&A)) -> Inspect<'r, A, Self> {
+        Inspect{iter: self, f: f}
     }
 
     /// An adaptation of an external iterator to the for-loop protocol of rust.
@@ -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,
@@ -1329,14 +1383,14 @@ impl<'self,
 
 /// An iterator that calls a function with a reference to each
 /// element before yielding it.
-pub struct Peek<'self, A, T> {
+pub struct Inspect<'self, A, T> {
     priv iter: T,
     priv f: &'self fn(&A)
 }
 
-impl<'self, A, T> Peek<'self, A, T> {
+impl<'self, A, T> Inspect<'self, A, T> {
     #[inline]
-    fn do_peek(&self, elt: Option<A>) -> Option<A> {
+    fn do_inspect(&self, elt: Option<A>) -> Option<A> {
         match elt {
             Some(ref a) => (self.f)(a),
             None => ()
@@ -1346,11 +1400,11 @@ impl<'self, A, T> Peek<'self, A, T> {
     }
 }
 
-impl<'self, A, T: Iterator<A>> Iterator<A> for Peek<'self, A, T> {
+impl<'self, A, T: Iterator<A>> Iterator<A> for Inspect<'self, A, T> {
     #[inline]
     fn next(&mut self) -> Option<A> {
         let next = self.iter.next();
-        self.do_peek(next)
+        self.do_inspect(next)
     }
 
     #[inline]
@@ -1359,15 +1413,17 @@ impl<'self, A, T: Iterator<A>> Iterator<A> for Peek<'self, A, T> {
     }
 }
 
-impl<'self, A, T: DoubleEndedIterator<A>> DoubleEndedIterator<A> for Peek<'self, A, T> {
+impl<'self, A, T: DoubleEndedIterator<A>> DoubleEndedIterator<A>
+for Inspect<'self, A, T> {
     #[inline]
     fn next_back(&mut self) -> Option<A> {
         let next = self.iter.next_back();
-        self.do_peek(next)
+        self.do_inspect(next)
     }
 }
 
-impl<'self, A, T: RandomAccessIterator<A>> RandomAccessIterator<A> for Peek<'self, A, T> {
+impl<'self, A, T: RandomAccessIterator<A>> RandomAccessIterator<A>
+for Inspect<'self, A, T> {
     #[inline]
     fn indexable(&self) -> uint {
         self.iter.indexable()
@@ -1375,7 +1431,7 @@ impl<'self, A, T: RandomAccessIterator<A>> RandomAccessIterator<A> for Peek<'sel
 
     #[inline]
     fn idx(&self, index: uint) -> Option<A> {
-        self.do_peek(self.iter.idx(index))
+        self.do_inspect(self.iter.idx(index))
     }
 }
 
@@ -1567,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];
@@ -1651,13 +1725,13 @@ mod tests {
     }
 
     #[test]
-    fn test_peek() {
+    fn test_inspect() {
         let xs = [1u, 2, 3, 4];
         let mut n = 0;
 
         let ys = xs.iter()
                    .map(|&x| x)
-                   .peek(|_| n += 1)
+                   .inspect(|_| n += 1)
                    .collect::<~[uint]>();
 
         assert_eq!(n, xs.len());
@@ -2011,11 +2085,11 @@ mod tests {
     }
 
     #[test]
-    fn test_random_access_peek() {
+    fn test_random_access_inspect() {
         let xs = [1, 2, 3, 4, 5];
 
-        // test .map and .peek that don't implement Clone
-        let it = xs.iter().peek(|_| {});
+        // test .map and .inspect that don't implement Clone
+        let it = xs.iter().inspect(|_| {});
         assert_eq!(xs.len(), it.indexable());
         for (i, elt) in xs.iter().enumerate() {
             assert_eq!(Some(elt), it.idx(i));
@@ -2027,7 +2101,6 @@ mod tests {
     fn test_random_access_map() {
         let xs = [1, 2, 3, 4, 5];
 
-        // test .map and .peek that don't implement Clone
         let it = xs.iter().map(|x| *x);
         assert_eq!(xs.len(), it.indexable());
         for (i, elt) in xs.iter().enumerate() {