Iterate mutably over &[1,2,3,4,5,6,7,8,9] to receive `&mut [1,4,7]` then `&mut [2,5,8]` then `&mut[3,6,9]`

Question

I'm trying to iterate over a slice broken into chunks, and return a tuple with the nth element of each chunk.

Example:

&[1,2,3,4,5,6,7,8,9]

I'd like to break this into chunks of size 3, and then iterate over the results, returning these tuples, one on each next() call:

&mut[1,4,7], &mut[2,5,8], &mut[3,6,9]

I know that for general stuff it isn't possible to return mutable stuff, mut this is clearly disjoint, and without unsafe code we can have the ChunksMut (https://doc.rust-lang.org/std/slice/struct.ChunksMut.html) iterator, so maybe there's a way!. For example, I can have 3 ChunksMut and then the compiler knows that the elements returned from them are disjoint.

This is my try for non mutable:

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=cfa7ca0bacbe6f1535050cd7dd5c537c

PS: I want to avoid Vec or any allocation on each iteration

Answer 1

so I always return a reference to its internal slice

The Iterator trait doesn't support this, because its contract allows the caller to extract several values and use all of them. For example, the following is permitted by Iterator but wouldn't be supported by your implementation:

// take two values out of the iterator
let a = it.next().unwrap();
let b = it.next().unwrap();

What you need is a "lending iterator" (also known as "streaming iterator"), see e.g. this crate. Writing lending iterators will become much easier once GATs are stabilized, but they still won't be supported by std::iter::Iterator.

Using the standard Iterator you can avoid allocation by using ArrayVec or equivalent replacement for Vec, as suggested by @Stargateur.

Answer 2

I'm pretty sure you wanted to get mutable references into the original slice using the iterator, resulting in &mut [&mut 1, &mut 4, &mut 7], &mut [&mut 2, &mut 5, &mut 8], &mut [&mut 3, &mut 6, &mut 9].

Without allocation / unsafe / external crates. Requires rust version 1.55 or greater:

fn iter_chunks<T, const CHUNK_SIZE: usize>(
    slice: &mut [T],
) -> impl Iterator<Item = [&mut T; CHUNK_SIZE]> + '_ {
    assert_eq!(slice.len() % CHUNK_SIZE, 0);
    let len = slice.len();
    let mut a: [_; CHUNK_SIZE] = array_collect(
        slice
            .chunks_mut(len / CHUNK_SIZE)
            .map(|iter| iter.iter_mut()),
    );
    (0..len / CHUNK_SIZE).map(move |_| array_collect(a.iter_mut().map(|i| i.next().unwrap())))
}

/// Builds an array from the first `N` items of an iterator
///
/// Panics:
///
/// If there are less then `N` items in the iterator
fn array_collect<T, const N: usize>(mut iter: impl Iterator<Item = T>) -> [T; N] {
    let a: [(); N] = [(); N];
    a.map(|_| iter.next().unwrap())
}

Without allocation, using an external crate:

We need to use arrayvec since Rust's array cannot be used with collect.

use arrayvec::ArrayVec;

fn main() {
    let slice = &mut [1, 2, 3, 4, 5, 6, 7, 8, 9];
    for (i, chunk) in iter_chunks::<_, 3>(slice).enumerate() {
        println!("{:?}", chunk);
        for t in chunk {
            *t = i;
        }
    }
    println!("slice: {:?}", slice);
}

fn iter_chunks<T, const CHUNK_SIZE: usize>(
    slice: &mut [T],
) -> impl Iterator<Item = ArrayVec<&mut T, CHUNK_SIZE>> + '_ {
    let len = slice.len();
    let mut a: ArrayVec<_, CHUNK_SIZE> = slice
        .chunks_mut(len / CHUNK_SIZE)
        .map(|chunk| chunk.iter_mut())
        .collect();
    (0..len / CHUNK_SIZE).map(move |_| {
        a.iter_mut()
            .map(|iter| iter.next().unwrap())
            .collect::<ArrayVec<_, CHUNK_SIZE>>()
    })
}

Output:

[1, 4, 7]
[2, 5, 8]
[3, 6, 9]
slice: [0, 1, 2, 0, 1, 2, 0, 1, 2]