Cycle a Rust iterator a given number of times
One simple way is to repeat the iterator itself, take the first 4 and flatten:
fn main() {
let v = vec![1, 2, 3];
let res = std::iter::repeat(v.iter())
.take(4)
.flatten()
.collect::<Vec<_>>();
dbg!(res);
}
Some micro-benchmark result using code in this gist comparing 3 different approaches:
- repeat-take-flatten in this answer
- hand-rolled loop
- a
cycle_n
implementation mimickingIterator::cycle
.
Kudos to rustc
, cycle_n
consistently outperforms the other two when the input is reasonably large whereas repeat-take-flatten performs the best for small input.
There is no such an iterator in the std lib.
If you know the iterator size, you can take
your number times the size of the iterator:
fn cycle_n_times<T: Clone>(slice: &[T], count: usize) -> impl Iterator<Item = &T> {
slice.iter().cycle().take(slice.len() * count)
}
Or you can write your own that is more general:
pub struct Ncycles<I> {
orig: I,
iter: I,
count: usize,
}
impl<I: Clone> Ncycles<I> {
fn new(iter: I, count: usize) -> Ncycles<I> {
Ncycles {
orig: iter.clone(),
iter,
count,
}
}
}
impl<I> Iterator for Ncycles<I>
where
I: Clone + Iterator,
{
type Item = <I as Iterator>::Item;
#[inline]
fn next(&mut self) -> Option<<I as Iterator>::Item> {
match self.iter.next() {
None if self.count == 0 => None,
None => {
self.iter = self.orig.clone();
self.count -= 1;
self.iter.next()
}
y => y,
}
}
}
#[test]
fn it_work() {
Ncycles::new(vec![1, 2, 3].iter(), 4).eq(&[1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3]);
}