Cycle a Rust iterator a given number of times

How do I cycle through an iterator a finite number of times?

I would expect the output of something like this to be 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3 and then stop:

vec![1, 2, 3].iter().cycle(4)
//                         ^ but .cycle() doesn't take an argument...

I don't know the length of the iterator to begin with.

Upvotes: 10

Views: 5324

Answers (3)

One could leverage slice::repeat, but I can't imagine this is very efficient:

let v = vec![1, 2, 3];
let it = v.iter();
println!("{:?}", it.map(|x| *x).collect::<Vec<i32>>().repeat(4).iter());

Upvotes: 0

edwardw
edwardw

Reputation: 13942

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 mimicking Iterator::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.

Micro-benchmark between repeat-take-flatten, hand-rolled loop and <code>cycle_n</code>

Upvotes: 15

Boiethios
Boiethios

Reputation: 42739

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]);
}

Upvotes: 4

Related Questions