Is there a clear syntax for iterating multiple iterators at once?

I need to iterate over 5 arrays at once. I'm converting this C code: (please forgive variable names, they're precalculated components of a math formula)

    for (int offset = 0; offset < width * height; offset++) {
        const double mu1_sq = mu1[offset]*mu1[offset];
        const double mu2_sq = mu2[offset]*mu2[offset];
        const double mu1_mu2 = mu1[offset]*mu2[offset];
        const double sigma1_sq = img1_sq_blur[offset] - mu1_sq;
        const double sigma2_sq = img2_sq_blur[offset] - mu2_sq;
        const double sigma12 = img1_img2_blur[offset] - mu1_mu2;

The Rust equivalent I came up with is this (the original C code used aliased pointers for mu2 and an output array, and Rust version makes it explicit):

let img_sq_iter = original.img_sq_blur.iter().cloned().zip(modified.img_sq_blur.iter().cloned());
for (img1_img2_blur, ((mu1, mut mu2_in_map_out),(img1_sq_blur, img2_sq_blur))) in 
     img1_img2_blur.iter().cloned().zip(original.mu.iter().cloned().zip(modified.mu.iter_mut()).zip(img_sq_iter)) {

My general issue with this that it's a long unreadable line with lots of zips and nested parenthesis.

All these vectors have compatible types, so if I make a mistake and the destructuring pattern won't have the same order of variables as the jumbled chain of iterators then I'll get wrong results, but the compiler may not notice.

Is there a better syntax in Rust for multiple iterators? Something where the iterator and name of variable that iterates it is not so disconnected?

Some imaginary syntax that I'd like:

for let zip({
    img1_img2_blur: img1_img2_blur.iter().cloned(),
    mu1: original.mu.iter().cloned(),
    mut mu2_in_map_out: modified.mu.iter_mut(),
    img1_sq_blur: original.img_sq_blur.iter().cloned(),
    img2_sq_blur: modified.img_sq_blur.iter().cloned(),
}) {
    … // ^ isn't that much easier to follow?
}
1 Like

Rust doesn't yet have the equivalent of variable length templates. But in the itertools (http://bluss.github.io/rust-itertools/doc/itertools/struct.Zip.html ) they offer the syntax:

Zip::new((it1, it2, it3, it4, it5))

Or this macro:
http://bluss.github.io/rust-itertools/doc/itertools/macro.izip!.html

That allows code like:

izip!(it1, it2, it3, it4, it5)

But in your code you can also iterate on a (0 .. N) interval and then index arrays... I think currently the Rust compiler doesn't remove the bound tests in this case.

2 Likes

You might just want to consider:

for offset in 0..width*height
{
}

Which maps very close to the original code....

1 Like

I'm afraid this wouldn't avoid bounds checks, as there's no way for the compiler to know that the vectors have the same size.

Thanks. I've switched to itertools::Zip, which is slightly better, although even in the macro version it still leaves variable names disconnected from iterators they iterate.

The state of the art in zipping performance is unfortunately the method mentioned in How to “zip” two slices efficiently; slice them to the same length and index iterate. Even the unsafe-using ZipSlices which was as good has regressed in recent rust and does not generate optimal or vectorizable code due to rustc/llvm codgen issues.

Anyway, here's a good rule: If you use the slice's length as loop counter bound, then llvm will remove bounds checks inside the loop. It's logical. You just have to make sure to use something that llvm sees is identical with the slice length.

I haven't verified that the trick works with three or more slices, probably does?

1 Like

The ZipSlices codegen issues are directly connected to llvm not trusting that the references are nonnull, or somehow losing that information. That means that ZipSlices can still generate the “perfect” (autovectorizable) loop for something that doesn't return references, like this:

struct CopySlice<'a, T: 'a>(&'a [T]);

unsafe impl<'a, T: Copy> itertools::misc::Slice for CopySlice<'a, T> {
    type Item = T;
    #[inline]
    fn len(&self) -> usize { self.0.len() }
    #[inline]
    unsafe fn get_unchecked(&mut self, i: usize) -> T { *self.0.get_unchecked(i) }
}