diff --git a/src/libcore/iterator.rs b/src/libcore/iterator.rs index 7ca40ae3a300..e233eb7feed6 100644 --- a/src/libcore/iterator.rs +++ b/src/libcore/iterator.rs @@ -18,6 +18,7 @@ pub trait Iterator { } pub trait IteratorUtil { + fn chain(self, other: Self) -> ChainIterator; fn zip>(self, other: U) -> ZipIterator; // FIXME: #5898: should be called map fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, Self>; @@ -31,6 +32,11 @@ pub trait IteratorUtil { } impl> IteratorUtil for T { + #[inline(always)] + fn chain(self, other: T) -> ChainIterator { + ChainIterator{a: self, b: other, flag: false} + } + #[inline(always)] fn zip>(self, other: U) -> ZipIterator { ZipIterator{a: self, b: other} @@ -86,6 +92,28 @@ impl> IteratorUtil for T { } } +pub struct ChainIterator { + priv a: T, + priv b: T, + priv flag: bool +} + +impl> Iterator for ChainIterator { + #[inline] + fn next(&mut self) -> Option { + 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 { priv a: T, priv b: U @@ -288,6 +316,20 @@ 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];