split_adaptor.rs 3.02 KB
Newer Older
1 2
use std::iter::Peekable;

3 4
use common::chunk::Chunk;

5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
/// Represents iterator adaptor that splits source iterator when supplied
/// predicate stops evaluating to true, producing iterator over iterators as
/// a result.
#[derive(Clone)]
pub struct Split<T, F>
where
    T: Clone + Iterator,
    T::Item: Clone, // required for Peekable
    F: Fn(&T::Item, &T::Item) -> bool,
{
    /// Original iterator to adapt, wrapped in `Peekable` to be able to call
    /// `self.predicate` on current + next elements.
    source: Peekable<T>,
    /// Predicate that returns 'false' when chunk has to end at the current
    /// element
    predicate: F,
}

impl<T, F> Iterator for Split<T, F>
where
    T: Clone + Iterator,
    T::Item: Clone,
    F: Fn(&T::Item, &T::Item) -> bool,
{
    type Item = Chunk<Peekable<T>>;

    fn next(&mut self) -> Option<Self::Item> {
        // Don't try creating chunks if underlying iterator is empty
        if self.source.peek().is_none() {
            return None;
        }

        // Save copy of the iterator to be returned with `Chunk` after original
        // one advanced to the next split point
        let chunk_source = self.source.clone();
        let mut chunk_count = 0;

        loop {
            // was already peeked previously, can't be None
            let a = self.source.next().unwrap();
            chunk_count += 1;

            match self.source.peek() {
                None => {
                    return Some(Chunk {
                        source: chunk_source,
                        count: chunk_count,
                    })
                }
                Some(x) => {
                    if !(self.predicate)(&a, x) {
                        continue;
                    } else {
                        return Some(Chunk {
                            source: chunk_source,
                            count: chunk_count,
                        });
                    }
                }
            }
        }
    }
}

/// Defines `split` method to all clonable iterators
pub trait SplitAdaptor: Iterator + Clone {
    /// Returns iterator adaptor that interprets original iterator as sequence
    /// of chunks with a border defined by supplied predicate.
    ///
    /// # Arguments
    /// * `predicate` - function/closure that controls when this iterator has to
    ///   be split into next chunk
    fn split<F>(&self, predicate: F) -> Split<Self, F>
    where
        Self::Item: Clone,
        F: Fn(&Self::Item, &Self::Item) -> bool,
    {
        Split {
            source: self.clone().peekable(),
            predicate: predicate,
        }
    }
}

impl<T> SplitAdaptor for T
where
    T: Clone + Iterator,
{
}

#[cfg(test)]
use itertools::assert_equal;

#[test]
fn test_split() {
    let data = vec![1, 2, 2, 2, 3, 3];

    let mut x = data.iter().split(|a, b| a != b);

    let x1 = x.next().unwrap();
    assert_equal(x1, [1].iter());
    let x2 = x.next().unwrap();
    assert_equal(x2, [2, 2, 2].iter());
    let x3 = x.next().unwrap();
    assert_equal(x3, [3, 3].iter());

    assert!(x.next().is_none());
}