chunks.rs

pub trait IteratorExt: Iterator + Sized {
    fn chunks(self, size: usize) -> Chunked<Self>;}

impl<T> IteratorExt for T
where T: Iterator
{
    fn chunks(self, size: usize) -> Chunked<Self> {
        chunk_iter(self, size)
    }
}

pub fn chunk_iter<I>(iter: I, size: usize) -> Chunked<I::IntoIter>
where
    I: IntoIterator
{
    Chunked::new(iter.into_iter(), size)
}

pub struct Chunked<I> {
    iter: I,
    size: usize
}

impl<I> Chunked<I> {
    fn new(iter: I, size: usize) -> Self {
        Chunked {iter, size}
    }
}

impl<I> Iterator for Chunked<I>
where
    I: Iterator,
{
    type Item = Vec<I::Item>;
    fn next(&mut self) -> Option<Self::Item> {
        let mut chunk = vec![];

        while chunk.len() < self.size {
            match self.iter.next() {
                Some(next) => chunk.push(next),
                None => break
            }
        }

        match chunk.len() {
            0 => None,
            _ => Some(chunk)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty() {
        assert_eq!(chunk_iter(std::iter::empty::<Vec<()>>(), 2).count(), 0)
    }

    #[test]
    fn one_element_two_chunk_iters() {
        let mut iter = chunk_iter(vec!["a"], 2);
        assert_eq!(iter.next(), Some(vec!["a"]));
    }

    #[test]
    fn two_element_two_chunk_iters() {
        let mut iter = chunk_iter(vec!["a", "b"], 2);
        assert_eq!(iter.next(), Some(vec!["a", "b"]));
    }

    #[test]
    fn three_element_two_chunk_iters() {
        let mut iter = chunk_iter(vec!["a", "b", "c"], 2);
        assert_eq!(iter.next(), Some(vec!["a", "b"]));
        assert_eq!(iter.next(), Some(vec!["c"]));
    }

    #[test]
    fn four_element_two_chunk_iters() {
        let mut iter = chunk_iter(vec!["a", "b", "c", "d"], 2);
        assert_eq!(iter.next(), Some(vec!["a", "b"]));
        assert_eq!(iter.next(), Some(vec!["c", "d"]));
    }

    #[test]
    fn four_element_two_chunk_ext() {
        assert_eq!(vec!["a", "b", "c", "d", "e"].into_iter().chunks(2).count(), 3)
    }
}