auto merge of #7604 : apasel422/rust/peek, r=huonw

This can be useful for inserting debugging code at different steps in an iterator pipeline.

1 files changed, 65 insertions, 0 deletions

diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index 8094b238749..4f5f1bd85c7 100644
--- a/src/libstd/iterator.rs
+++ b/src/libstd/iterator.rs
@@ -245,6 +245,25 @@ pub trait IteratorUtil<A> {
     fn flat_map_<'r, B, U: Iterator<B>>(self, f: &'r fn(A) -> U)
         -> FlatMapIterator<'r, A, B, Self, U>;
 
+    /// Creates an iterator that calls a function with a reference to each
+    /// element before yielding it. This is often useful for debugging an
+    /// iterator pipeline.
+    ///
+    /// # Example
+    ///
+    /// ~~~ {.rust}
+    ///let xs = [1u, 4, 2, 3, 8, 9, 6];
+    ///let sum = xs.iter()
+    ///            .transform(|&x| x)
+    ///            .peek_(|&x| debug!("filtering %u", x))
+    ///            .filter(|&x| x % 2 == 0)
+    ///            .peek_(|&x| debug!("%u made it through", x))
+    ///            .sum();
+    ///println(sum.to_str());
+    /// ~~~
+    // FIXME: #5898: should be called `peek`
+    fn peek_<'r>(self, f: &'r fn(&A)) -> PeekIterator<'r, A, Self>;
+
     /// An adaptation of an external iterator to the for-loop protocol of rust.
     ///
     /// # Example
@@ -442,6 +461,12 @@ impl<A, T: Iterator<A>> IteratorUtil<A> for T {
         FlatMapIterator{iter: self, f: f, subiter: None }
     }
 
+    // FIXME: #5898: should be called `peek`
+    #[inline]
+    fn peek_<'r>(self, f: &'r fn(&A)) -> PeekIterator<'r, A, T> {
+        PeekIterator{iter: self, f: f}
+    }
+
     /// A shim implementing the `for` loop iteration protocol for iterator objects
     #[inline]
     fn advance(&mut self, f: &fn(A) -> bool) -> bool {
@@ -1041,6 +1066,32 @@ impl<'self, A, T: Iterator<A>, B, U: Iterator<B>> Iterator<B> for
     }
 }
 
+/// An iterator that calls a function with a reference to each
+/// element before yielding it.
+pub struct PeekIterator<'self, A, T> {
+    priv iter: T,
+    priv f: &'self fn(&A)
+}
+
+impl<'self, A, T: Iterator<A>> Iterator<A> for PeekIterator<'self, A, T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        let next = self.iter.next();
+
+        match next {
+            Some(ref a) => (self.f)(a),
+            None => ()
+        }
+
+        next
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        self.iter.size_hint()
+    }
+}
+
 /// An iterator which just modifies the contained state throughout iteration.
 pub struct UnfoldrIterator<'self, A, St> {
     priv f: &'self fn(&mut St) -> Option<A>,
@@ -1237,6 +1288,20 @@ mod tests {
     }
 
     #[test]
+    fn test_peek() {
+        let xs = [1u, 2, 3, 4];
+        let mut n = 0;
+
+        let ys = xs.iter()
+                   .transform(|&x| x)
+                   .peek_(|_| n += 1)
+                   .collect::<~[uint]>();
+
+        assert_eq!(n, xs.len());
+        assert_eq!(xs, ys.as_slice());
+    }
+
+    #[test]
     fn test_unfoldr() {
         fn count(st: &mut uint) -> Option<uint> {
             if *st < 10 {