auto merge of #5955 : thestinger/rust/iterator, r=graydon

2 files changed, 251 insertions, 61 deletions

diff --git a/src/libcore/iterator.rs b/src/libcore/iterator.rs
index 8bd6c73fc7d..4929b1b8dba 100644
--- a/src/libcore/iterator.rs
+++ b/src/libcore/iterator.rs
@@ -8,7 +8,14 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-//! Composable iterator objects
+/*! Composable external iterators
+
+The `Iterator` trait defines an interface for objects which implement iteration as a state machine.
+
+Algorithms like `zip` are provided as `Iterator` implementations which wrap other objects
+implementing the `Iterator` trait.
+
+*/
 
 use prelude::*;
 
@@ -17,19 +24,35 @@ pub trait Iterator<A> {
     fn next(&mut self) -> Option<A>;
 }
 
+/// Iterator adaptors provided for every `Iterator` implementation. The adaptor objects are also
+/// implementations of the `Iterator` trait.
+///
+/// In the future these will be default methods instead of a utility trait.
 pub trait IteratorUtil<A> {
+    fn chain(self, other: Self) -> ChainIterator<Self>;
     fn zip<B, U: Iterator<B>>(self, other: U) -> ZipIterator<Self, U>;
     // 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 skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, Self>;
+    fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, Self>;
+    fn skip(self, n: uint) -> SkipIterator<Self>;
+    fn take(self, n: uint) -> TakeIterator<Self>;
     fn advance(&mut self, f: &fn(A) -> bool);
 }
 
+/// Iterator adaptors provided for every `Iterator` implementation. The adaptor objects are also
+/// implementations of the `Iterator` trait.
+///
+/// In the future these will be default methods instead of a utility trait.
 impl<A, T: Iterator<A>> IteratorUtil<A> for T {
     #[inline(always)]
+    fn chain(self, other: T) -> ChainIterator<T> {
+        ChainIterator{a: self, b: other, flag: false}
+    }
+
+    #[inline(always)]
     fn zip<B, U: Iterator<B>>(self, other: U) -> ZipIterator<T, U> {
         ZipIterator{a: self, b: other}
     }
@@ -51,15 +74,25 @@ impl<A, T: Iterator<A>> IteratorUtil<A> for T {
     }
 
     #[inline(always)]
-    fn dropwhile<'r>(self, predicate: &'r fn(&A) -> bool) -> DropWhileIterator<'r, A, T> {
-        DropWhileIterator{iter: self, flag: false, predicate: predicate}
+    fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, T> {
+        SkipWhileIterator{iter: self, flag: false, predicate: predicate}
     }
 
     #[inline(always)]
-    fn takewhile<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, T> {
+    fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, T> {
         TakeWhileIterator{iter: self, flag: false, predicate: predicate}
     }
 
+    #[inline(always)]
+    fn skip(self, n: uint) -> SkipIterator<T> {
+        SkipIterator{iter: self, n: n}
+    }
+
+    #[inline(always)]
+    fn take(self, n: uint) -> TakeIterator<T> {
+        TakeIterator{iter: self, n: n}
+    }
+
     /// A shim implementing the `for` loop iteration protocol for iterator objects
     #[inline]
     fn advance(&mut self, f: &fn(A) -> bool) {
@@ -74,6 +107,28 @@ impl<A, T: Iterator<A>> IteratorUtil<A> for T {
     }
 }
 
+pub struct ChainIterator<T> {
+    priv a: T,
+    priv b: T,
+    priv flag: bool
+}
+
+impl<A, T: Iterator<A>> Iterator<A> for ChainIterator<T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        if self.flag {
+            self.b.next()
+        } else {
+            match self.a.next() {
+                Some(x) => return Some(x),
+                _ => ()
+            }
+            self.flag = true;
+            self.b.next()
+        }
+    }
+}
+
 pub struct ZipIterator<T, U> {
     priv a: T,
     priv b: U
@@ -89,6 +144,21 @@ impl<A, B, T: Iterator<A>, U: Iterator<B>> Iterator<(A, B)> for ZipIterator<T, U
     }
 }
 
+pub struct MapIterator<'self, A, B, T> {
+    priv iter: T,
+    priv f: &'self fn(A) -> B
+}
+
+impl<'self, A, B, T: Iterator<A>> Iterator<B> for MapIterator<'self, A, B, T> {
+    #[inline]
+    fn next(&mut self) -> Option<B> {
+        match self.iter.next() {
+            Some(a) => Some((self.f)(a)),
+            _ => None
+        }
+    }
+}
+
 pub struct FilterIterator<'self, A, T> {
     priv iter: T,
     priv predicate: &'self fn(&A) -> bool
@@ -108,21 +178,6 @@ impl<'self, A, T: Iterator<A>> Iterator<A> for FilterIterator<'self, A, T> {
     }
 }
 
-pub struct MapIterator<'self, A, B, T> {
-    priv iter: T,
-    priv f: &'self fn(A) -> B
-}
-
-impl<'self, A, B, T: Iterator<A>> Iterator<B> for MapIterator<'self, A, B, T> {
-    #[inline]
-    fn next(&mut self) -> Option<B> {
-        match self.iter.next() {
-            Some(a) => Some((self.f)(a)),
-            _ => None
-        }
-    }
-}
-
 pub struct EnumerateIterator<T> {
     priv iter: T,
     priv count: uint
@@ -142,13 +197,13 @@ impl<A, T: Iterator<A>> Iterator<(uint, A)> for EnumerateIterator<T> {
     }
 }
 
-pub struct DropWhileIterator<'self, A, T> {
+pub struct SkipWhileIterator<'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> {
+impl<'self, A, T: Iterator<A>> Iterator<A> for SkipWhileIterator<'self, A, T> {
     #[inline]
     fn next(&mut self) -> Option<A> {
         let mut next = self.iter.next();
@@ -199,3 +254,176 @@ impl<'self, A, T: Iterator<A>> Iterator<A> for TakeWhileIterator<'self, A, T> {
         }
     }
 }
+
+pub struct SkipIterator<T> {
+    priv iter: T,
+    priv n: uint
+}
+
+impl<A, T: Iterator<A>> Iterator<A> for SkipIterator<T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        let mut next = self.iter.next();
+        if self.n == 0 {
+            next
+        } else {
+            let n = self.n;
+            for n.times {
+                match next {
+                    Some(_) => {
+                        next = self.iter.next();
+                        loop
+                    }
+                    None => {
+                        self.n = 0;
+                        return None
+                    }
+                }
+            }
+            self.n = 0;
+            next
+        }
+    }
+}
+
+pub struct TakeIterator<T> {
+    priv iter: T,
+    priv n: uint
+}
+
+impl<A, T: Iterator<A>> Iterator<A> for TakeIterator<T> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        let next = self.iter.next();
+        if self.n != 0 {
+            self.n -= 1;
+            next
+        } else {
+            None
+        }
+    }
+}
+
+pub struct UnfoldrIterator<'self, A, St> {
+    priv f: &'self fn(&mut St) -> Option<A>,
+    priv state: St
+}
+
+pub impl<'self, A, St> UnfoldrIterator<'self, A, St> {
+    #[inline]
+    fn new(f: &'self fn(&mut St) -> Option<A>, initial_state: St) -> UnfoldrIterator<'self, A, St> {
+        UnfoldrIterator {
+            f: f,
+            state: initial_state
+        }
+    }
+}
+
+impl<'self, A, St> Iterator<A> for UnfoldrIterator<'self, A, St> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        (self.f)(&mut self.state)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use prelude::*;
+
+    #[test]
+    fn test_iterator_chain() {
+        let xs = [0u, 1, 2, 3, 4, 5];
+        let ys = [30, 40, 50, 60];
+        let expected = [0, 1, 2, 3, 4, 5, 30, 40, 50, 60];
+        let mut it = xs.iter().chain(ys.iter());
+        let mut i = 0;
+        for it.advance |&x: &uint| {
+            assert_eq!(x, expected[i]);
+            i += 1;
+        }
+        assert_eq!(i, expected.len());
+    }
+
+    #[test]
+    fn test_iterator_enumerate() {
+        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);
+        }
+    }
+
+    #[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];
+        let mut it = xs.iter().take_while(|&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_skip_while() {
+        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
+        let ys = [15, 16, 17, 19];
+        let mut it = xs.iter().skip_while(|&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_skip() {
+        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
+        let ys = [13, 15, 16, 17, 19, 20, 30];
+        let mut it = xs.iter().skip(5);
+        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_take() {
+        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
+        let ys = [0u, 1, 2, 3, 5];
+        let mut it = xs.iter().take(5);
+        let mut i = 0;
+        for it.advance |&x: &uint| {
+            assert_eq!(x, ys[i]);
+            i += 1;
+        }
+        assert_eq!(i, ys.len());
+    }
+
+    #[test]
+    fn test_unfoldr() {
+        fn count(st: &mut uint) -> Option<uint> {
+            if *st < 10 {
+                let ret = Some(*st);
+                *st += 1;
+                ret
+            } else {
+                None
+            }
+        }
+
+        let mut it = UnfoldrIterator::new(count, 0);
+        let mut i = 0;
+        for it.advance |counted| {
+            assert_eq!(counted, i);
+            i += 1;
+        }
+        assert_eq!(i, 10);
+    }
+}
diff --git a/src/libcore/vec.rs b/src/libcore/vec.rs
index 45cc9618f59..139fcedad27 100644
--- a/src/libcore/vec.rs
+++ b/src/libcore/vec.rs
@@ -4474,42 +4474,4 @@ mod tests {
             i += 1;
         }
     }
-
-    #[test]
-    fn test_iterator_enumerate() {
-        use iterator::*;
-        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);
-        }
-    }
-
-    #[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());
-    }
 }