How to obtain the chunk index in Rayon's par_chunks_mut

Question

I have some data and I want to process it and use it to fill an array that already exists. For example suppose I want to repeat each value 4 times (playground):

use rayon::prelude::*; // 1.3.0

fn main() {
    let input = vec![4, 7, 2, 3, 5, 8];

    // This already exists.
    let mut output = vec![0; input.len() * 4];

    output.par_chunks_mut(4).for_each(|slice| {
        for x in slice.iter_mut() {
            *x = input[?];
        }
    });
}

This almost works but Rayon doesn't pass the chunk index to me so I can't put anything in input[?]. Is there an efficient solution?

Answer 1

The easiest thing to do is avoid the need for an index at all. For this example, we can just zip the iterators:

use rayon::prelude::*; // 1.3.0

fn main() {
    let input = vec![4, 7, 2, 3, 5, 8];
    let mut output = vec![0; input.len() * 4];

    // Can also use `.zip(&input)` if you don't want to give up ownership
    output.par_chunks_mut(4).zip(input).for_each(|(o, i)| {
        for o in o {
            *o = i
        }
    });

    println!("{:?}", output)
}

For traditional iterators, this style of implementation is beneficial as it avoids unneeded bounds checks which would otherwise be handled by the iterator. I'm not sure that Rayon benefits from the exact same thing, but I also don't see any reason it wouldn't.

Answer 2

Rayon provides an enumerate() function for most of its iterators that works just like the non-parallel counterpart:

let input = vec![4, 7, 2, 3, 5, 8];
let mut output = vec![0; input.len() * 4];

output.par_chunks_mut(4).enumerate().for_each(|(i, slice)| {
    for x in slice.iter_mut() {
        *x = input[i];
    }
});