0
\$\begingroup\$

repeat_each transforms an iterator by repeating each value count times. It should be equivalent to .flat_map(|x| iter::repeat(x).take(count)).

#![cfg_attr(feature = "override-try-fold", feature(try_trait))]
#![cfg_attr(feature = "test-bench", feature(test))]

use std::iter::{ExactSizeIterator, FusedIterator};
#[cfg(feature = "override-try-fold")]
use std::ops::Try;

pub trait IteratorExt: Iterator {
    /// Repeats each value in `self` `count` times
    fn repeat_each(self, count: usize) -> RepeatEach<Self>
    where
        Self::Item: Clone,
        Self: Sized,
    {
        RepeatEach {
            count,
            idx: 0,
            value: None,
            inner: self,
        }
    }
}
impl<T: Iterator> IteratorExt for T {}

pub struct RepeatEach<I: Iterator>
where
    I::Item: Clone,
{
    count: usize,
    idx: usize,
    value: Option<I::Item>,
    inner: I,
}
impl<I: Iterator> RepeatEach<I>
where
    I::Item: Clone,
{
    #[inline]
    fn transform_size(&self, inner_size: usize) -> usize {
        inner_size * self.count
            + match self.value {
                Some(_) => self.count - self.idx,
                None => 0,
            }
    }
}
impl<I: Iterator> Iterator for RepeatEach<I>
where
    I::Item: Clone,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<I::Item> {
        if self.count == 0 {
            return None;
        }
        if self.idx == 0 && self.value.is_none() {
            self.value = self.inner.next();
        }
        if self.idx == self.count {
            self.value = self.inner.next();
            self.idx = 0;
        }
        self.idx += 1;
        self.value.clone()
    }

    #[cfg(feature = "override-try-fold")]
    fn try_fold<B, F, R>(&mut self, mut init: B, mut f: F) -> R
    where
        F: FnMut(B, Self::Item) -> R,
        R: Try<Ok = B>,
    {
        if let Some(x) = &self.value {
            for i in self.idx..self.count {
                self.idx = i + 1;
                init = f(init, x.clone())?;
            }
        }
        while let Some(x) = self.inner.next() {
            self.value = Some(x.clone());
            for i in 0..self.count {
                self.idx = i + 1;
                init = f(init, x.clone())?;
            }
        }
        R::from_ok(init)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (inner_min, inner_max) = self.inner.size_hint();
        (
            self.transform_size(inner_min),
            inner_max.map(|x| self.transform_size(x)),
        )
    }
}
impl<I: FusedIterator> FusedIterator for RepeatEach<I> where I::Item: Clone {}
impl<I: ExactSizeIterator> ExactSizeIterator for RepeatEach<I>
where
    I::Item: Clone,
{
    fn len(&self) -> usize {
        self.transform_size(self.inner.len())
    }
}

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

    #[test]
    fn next() {
        assert_eq!(
            ([] as [u8; 0])
                .iter()
                .copied()
                .repeat_each(0)
                .collect::<Vec<_>>()[..],
            []
        );
        assert_eq!(
            ([] as [u8; 0])
                .iter()
                .copied()
                .repeat_each(1)
                .collect::<Vec<_>>()[..],
            []
        );
        assert_eq!(
            ([] as [u8; 0])
                .iter()
                .copied()
                .repeat_each(2)
                .collect::<Vec<_>>()[..],
            []
        );
        assert_eq!(
            [1_u8].iter().copied().repeat_each(0).collect::<Vec<_>>()[..],
            []
        );
        assert_eq!(
            [1_u8].iter().copied().repeat_each(1).collect::<Vec<_>>()[..],
            [1]
        );
        assert_eq!(
            [1_u8].iter().copied().repeat_each(2).collect::<Vec<_>>()[..],
            [1, 1]
        );
        assert_eq!(
            [1_u8, 2].iter().copied().repeat_each(0).collect::<Vec<_>>()[..],
            []
        );
        assert_eq!(
            [1_u8, 2].iter().copied().repeat_each(1).collect::<Vec<_>>()[..],
            [1, 2]
        );
        assert_eq!(
            [1_u8, 2].iter().copied().repeat_each(2).collect::<Vec<_>>()[..],
            [1, 1, 2, 2]
        );
    }

    #[test]
    fn size_hint_new() {
        assert_eq!(
            ([] as [u8; 0]).iter().repeat_each(0).size_hint(),
            (0, Some(0))
        );
        assert_eq!(
            ([] as [u8; 0]).iter().repeat_each(1).size_hint(),
            (0, Some(0))
        );
        assert_eq!(
            ([] as [u8; 0]).iter().repeat_each(2).size_hint(),
            (0, Some(0))
        );
        assert_eq!([1_u8].iter().repeat_each(0).size_hint(), (0, Some(0)));
        assert_eq!([1_u8].iter().repeat_each(1).size_hint(), (1, Some(1)));
        assert_eq!([1_u8].iter().repeat_each(2).size_hint(), (2, Some(2)));
        assert_eq!([1_u8, 2].iter().repeat_each(0).size_hint(), (0, Some(0)));
        assert_eq!([1_u8, 2].iter().repeat_each(1).size_hint(), (2, Some(2)));
        assert_eq!([1_u8, 2].iter().repeat_each(2).size_hint(), (4, Some(4)));
    }

    #[test]
    fn size_hint_mut() {
        for &(original, count) in
            &([(&[1], 1), (&[1], 2), (&[1, 2], 1), (&[1, 2], 2)] as [(&[u8], usize); 4])
        {
            let mut len = original.len() * count;
            let mut iter = original.iter().repeat_each(count);
            assert_eq!(iter.size_hint(), (len, Some(len)));
            while let Some(_) = iter.next() {
                len -= 1;
                assert_eq!(iter.size_hint(), (len, Some(len)));
            }
            assert_eq!(iter.size_hint(), (0, Some(0)));
        }
    }

    #[test]
    fn exact_size_iterator_len_new() {
        assert_eq!(([] as [u8; 0]).iter().repeat_each(0).len(), 0);
        assert_eq!(([] as [u8; 0]).iter().repeat_each(1).len(), 0);
        assert_eq!(([] as [u8; 0]).iter().repeat_each(2).len(), 0);
        assert_eq!([1_u8].iter().repeat_each(0).len(), 0);
        assert_eq!([1_u8].iter().repeat_each(1).len(), 1);
        assert_eq!([1_u8].iter().repeat_each(2).len(), 2);
        assert_eq!([1_u8, 2].iter().repeat_each(0).len(), 0);
        assert_eq!([1_u8, 2].iter().repeat_each(1).len(), 2);
        assert_eq!([1_u8, 2].iter().repeat_each(2).len(), 4);
    }

    #[test]
    fn exact_size_iterator_len_mut() {
        for &(original, count) in
            &([(&[1], 1), (&[1], 2), (&[1, 2], 1), (&[1, 2], 2)] as [(&[u8], usize); 4])
        {
            let mut len = original.len() * count;
            let mut iter = original.iter().repeat_each(count);
            assert_eq!(iter.len(), len);
            while let Some(_) = iter.next() {
                len -= 1;
                assert_eq!(iter.len(), len);
            }
            assert_eq!(iter.len(), 0);
        }
    }

    #[test]
    fn test_try_fold_ok() {
        let mut it = [75, 10].iter().repeat_each(3);
        let sum = it.try_fold(0u8, |acc, &x| acc.checked_add(x));
        assert_eq!(sum, Some(255));
        assert_eq!(it.len(), 0);
        assert_eq!(it.next(), None);
    }

    #[test]
    fn test_try_fold_error() {
        let mut it = [75, 20, 10].iter().repeat_each(3);
        let sum = it.try_fold(0u8, |acc, &x| acc.checked_add(x));
        assert_eq!(sum, None);
        assert_eq!(it.len(), 4);
        assert_eq!(it.next(), Some(&20));
        assert_eq!(it.next(), Some(&10));
        assert_eq!(it.next(), Some(&10));
        assert_eq!(it.next(), Some(&10));
        assert_eq!(it.next(), None);
    }
}

#[cfg(test)]
#[cfg(feature = "test-bench")]
mod benches {   
    extern crate test;

    use super::IteratorExt;
    use std::iter;
    use test::Bencher;

    #[bench]
    fn bench_repeat_each(b: &mut Bencher) {
        b.iter(|| [1_u8, 2, 3, 4, 5].iter().copied().repeat_each(4).collect::<Vec<_>>());
    }

    #[bench]
    fn bench_flat_map_repeat_take(b: &mut Bencher) {
        b.iter(|| [1_u8, 2, 3, 4, 5].iter().copied().flat_map(|x| iter::repeat(x).take(4)).collect::<Vec<_>>());
    }
}

Benchmark results:

test benches::bench_flat_map_repeat_take ... bench:         629 ns/iter (+/- 111)
test benches::bench_repeat_each          ... bench:         178 ns/iter (+/- 114)

Is it possible to optimize the time or memory usage of this? Are there any other improvements I could make?

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.