Reorder iteration without heap allocation

what are my options when I want to iterate in a non-consecutive fashion? e.g. I have &[1,2,3,4,5,6,7,8,9,10,11,12] and want an iterator that yields 1,2,3, 7,8,9, 5,6,7, 10,11,12.

The reordering I'd like to achieve is not completely arbitrary, but limited to shuffling equally sized chunks around. However that doesn't seem to help much, since a chunks-iterator still yields items in the original order.

The tricky part: I'd like to avoid heap allocations (putting the iterators in a Vec and shuffling them around there or so) for no_std and performance reasons.

I want to be generic over array size and chunk size, both are known at compile time.

Some background to illustrate: the concrete problem I'm trying to solve is pixel reordering for several adjacent LED matrix displays.
since the displays are daisy chained there is a discrepancy between visual & backing store pixel order:

visualy, one wants to just store, process & show rows of pixels:

|1 2 3|  |4 5 6|  |7 8 9|
|a b c|  |d e f|  |g h i|
  M1        M2       M3

however, if sent to the chain of displays in that order, things would get garbled because of the the connection topology (first, all pixels of M1, then all of M2, then M3):

|1 2 3| /|7 8 9| /|d e f|
|4 5 6|/ |a b c|/ |g h i|
  M1        M2       M3

I don't think the general problem is possible to solve without allocating some scratch space or destroying performance by doing loads of redundant comparisons/reads. Especially if all you've got to work with is a &[T] slice.

One option is to create an iterator which scans your slice every time you call next(), keeping track of the last chunk it yielded and looking for the next "smallest" chunk. That's O(n2), though.

Assuming your source is always a contiguous slice, it isn't too difficult:

fn array_chunks_reordered<T, const C: usize>(
    src: &[T],
    order: impl IntoIterator<Item = usize>,
) -> impl Iterator<Item = &[T; C]> {
    order.into_iter().map(|chunk| {
        let start = chunk.checked_mul(C).unwrap();
        src[start..start + C].try_into().unwrap()
    })
}

fn main() {
    let src = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
    let order = [0, 2, 1, 3];
    let reordered: Vec<&[_; 3]> = array_chunks_reordered(src, order).collect();
    println!("{reordered:?}");
}

Annoyingly, the chunk size can't be explicitly set with ::<_, C>, due to order being an impl IntoIterator argument, but you can change the type of order to match whatever variable it is coming from.

Then don't make it one! An ordinary generic works just fine:

fn array_chunks_reordered<const C: usize, T, I>(src: &[T], order: I) -> impl Iterator<Item = &[T; C]>
where
    I: IntoIterator<Item = usize>
{
    order.into_iter().map(|chunk| {
        let start = chunk.checked_mul(C).unwrap();
        src[start..start + C].try_into().unwrap()
    })
}
5 Likes

works like a charm, thanks!