如何将此循环转换为惯用的Rust？

Question

Here is a working Rust function:

fn foo(src: &[u8]) -> Vec<u8> {
    let dst_len = (src.len() / 3) * 4;
    let mut dst = vec![0 as u8; dst_len];

    let mut si = 0;
    let mut di = 0;
    let n = (src.len() / 3) * 3;
    for _ in (0 .. n).step_by(3) {
        let v = bar(src[si], src[si+1], src[si+2]);
        dst[di+0] = baz(v, 0);
        dst[di+1] = baz(v, 1);
        dst[di+2] = baz(v, 2);
        dst[di+3] = baz(v, 3);
        si += 3;
        di += 4;
    }

    dst
}

It works but that loop doesn't seem like idiomatic Rust. It indexes into arrays using manually managed indices, pretty much like a for loop in C.

Is there a way to achieve the same result using Rust iterators? I think chunked_exact would work for iterating over src , but what about dst ? What iterator could I zip with src.chunked_exact to write into dst in chunks?

Answer 1

What is "idiomatic" can be a matter of opinion, but you can make use of more iterator methods, like zip and chunks_exact_mut :

fn foo(src: &[u8]) -> Vec<u8> {
    let dst_len = (src.len() / 3) * 4;
    let mut dst = vec![0 as u8; dst_len];
    for (s, d) in src.chunks_exact(3).zip(dst.chunks_exact_mut(4)) {
        let v = bar(s[0], s[1], s[2]);
        d[0] = baz(v, 0);
        d[1] = baz(v, 1);
        d[2] = baz(v, 2);
        d[3] = baz(v, 3);
    }
    dst
}

I used chunks_exact and chunks_exact_mut rather than chunks because it guarantees that the slice has the requested length, making them available separately if you need them. This seems to match your original code, which rounds off the length to an exact number of steps.

Answer 2

I would go with that:

fn foo(src: &[u8]) -> Vec<u8> {
    src.chunks_exact(3)
        .map(|s| bar(s[0], s[1], s[2]))
        .flat_map(|v| (0..4).map(move |i| baz(v, i)))
        .collect()
}

Or if you prefer:

fn foo(src: &[u8]) -> Vec<u8> {
    src.chunks_exact(3)
        .map(|s| bar(s[0], s[1], s[2]))
        .flat_map(|v| vec![baz(v, 0), baz(v, 1), baz(v, 2), baz(v, 3)])
        .collect()
}

Not sure if you find one of them better.