add dropwhile and takewhile iterators

Author: thestinger

src/libcore/iterator.rs

@@ -22,6 +22,8 @@ pub trait IteratorUtil<A> {
     // FIXME: #5898: should be called map
     fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, Self>;
     fn filter<'r>(self, predicate: &'r fn(&A) -> bool) -> FilterIterator<'r, A, Self>;
+    fn dropwhile<'r>(self, predicate: &'r fn(&A) -> bool) -> DropWhileIterator<'r, A, Self>;
+    fn takewhile<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, Self>;
     fn enumerate(self) -> EnumerateIterator<Self>;
     fn advance(&mut self, f: &fn(A) -> bool);
 }
@@ -48,6 +50,16 @@ impl<A, T: Iterator<A>> IteratorUtil<A> for T {
         EnumerateIterator{iter: self, count: 0}
     }
 
+    #[inline(always)]
+    fn dropwhile<'r>(self, predicate: &'r fn(&A) -> bool) -> DropWhileIterator<'r, A, T> {
+        DropWhileIterator{iter: self, flag: false, predicate: predicate}
+    }
+
+    #[inline(always)]
+    fn takewhile<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, T> {
+        TakeWhileIterator{iter: self, flag: false, predicate: predicate}
+    }
+
     /// A shim implementing the `for` loop iteration protocol for iterator objects
     #[inline]
     fn advance(&mut self, f: &fn(A) -> bool) {
@@ -129,3 +141,61 @@ impl<A, T: Iterator<A>> Iterator<(uint, A)> for EnumerateIterator<T> {
         }
     }
 }
+
+pub struct DropWhileIterator<'self, A, T> {
+    priv iter: T,
+    priv flag: bool,
+    priv predicate: &'self fn(&A) -> bool
+}
+
+impl<'self, A, T: Iterator<A>> Iterator<A> for DropWhileIterator<'self, A, T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        let mut next = self.iter.next();
+        if self.flag {
+            next
+        } else {
+            loop {
+                match next {
+                    Some(x) => {
+                        if (self.predicate)(&x) {
+                            next = self.iter.next();
+                            loop
+                        } else {
+                            self.flag = true;
+                            return Some(x)
+                        }
+                    }
+                    None => return None
+                }
+            }
+        }
+    }
+}
+
+pub struct TakeWhileIterator<'self, A, T> {
+    priv iter: T,
+    priv flag: bool,
+    priv predicate: &'self fn(&A) -> bool
+}
+
+impl<'self, A, T: Iterator<A>> Iterator<A> for TakeWhileIterator<'self, A, T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        if self.flag {
+            None
+        } else {
+            match self.iter.next() {
+                Some(x) => {
+                    if (self.predicate)(&x) {
+                        Some(x)
+                    } else {
+                        self.flag = true;
+                        None
+                    }
+                }
+                None => None
+            }
+        }
+    }
+}

src/libcore/vec.rs

@@ -4478,18 +4478,38 @@ mod tests {
     #[test]
     fn test_iterator_enumerate() {
         use iterator::*;
-        let xs = [0u,1,2,3,4,5];
+        let xs = [0u, 1, 2, 3, 4, 5];
         let mut it = xs.iter().enumerate();
         for it.advance |(i, &x): (uint, &uint)| {
             assert_eq!(i, x);
         }
     }
-}
 
-// Local Variables:
-// mode: rust;
-// fill-column: 78;
-// indent-tabs-mode: nil
-// c-basic-offset: 4
-// buffer-file-coding-system: utf-8-unix
-// End:
+    #[test]
+    fn test_iterator_takewhile() {
+        use iterator::*;
+        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
+        let ys = [0u, 1, 2, 3, 5, 13];
+        let mut it = xs.iter().takewhile(|&x| *x < 15u);
+        let mut i = 0;
+        for it.advance |&x: &uint| {
+            assert_eq!(x, ys[i]);
+            i += 1;
+        }
+        assert_eq!(i, ys.len());
+    }
+
+    #[test]
+    fn test_iterator_dropwhile() {
+        use iterator::*;
+        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
+        let ys = [15, 16, 17, 19];
+        let mut it = xs.iter().dropwhile(|&x| *x < 15u);
+        let mut i = 0;
+        for it.advance |&x: &uint| {
+            assert_eq!(x, ys[i]);
+            i += 1;
+        }
+        assert_eq!(i, ys.len());
+    }
+}