How to turn an array of iterators into an iterator of arrays?

Question

Let's say I have the following array of iterators:

let arr = [0..9, 2..8, 4..8, 3..5];

How do I create an iterator that takes one element from each array every time .next() is called and returns a new array with these contents? I looked at itertools::izip and itertools::multizip, but these all didn't work with an array.

In this case, if such a function were called zip_array, I would expect the following result:

for [a, b, c, d] in zip_array(arr) {
    println!("{a} {b} {c} {d}");
}
println!("end of iteration");
// 0 2 4 3
// 1 3 5 4
// end of iteration

How could one implement this function?

Chayim Friedman · Accepted Answer · 2022-12-13 11:11:14Z

On nightly, this is very simple to implement (but note that the nightly implementation below may be faster):

#![feature(array_methods, array_try_map)]

pub struct ZipArray<T, const N: usize> {
    array: [T; N],
}

pub fn zip_array<T: Iterator, const N: usize>(array: [T; N]) -> ZipArray<T, N> {
    ZipArray { array }
}

impl<T: Iterator, const N: usize> Iterator for ZipArray<T, N> {
    type Item = [T::Item; N];

    fn next(&mut self) -> Option<Self::Item> {
        self.array.each_mut().try_map(|i| i.next())
    }
}

On stable, this require a little more work and unsafe code:

impl<T: Iterator, const N: usize> Iterator for ZipArray<T, N> {
    type Item = [T::Item; N];

    fn next(&mut self) -> Option<Self::Item> {
        // SAFETY: It is always valid to `assume_init()` an array of `MaybeUninit`s (can be replaced
        // with `MaybeUninit::uninit_array()` once stable).
        let mut result: [MaybeUninit<T::Item>; N] = unsafe { MaybeUninit::uninit().assume_init() };
        for (item, iterator) in std::iter::zip(&mut result, &mut self.array) {
            item.write(iterator.next()?);
        }
        // SAFETY: We initialized the array above (can be replaced with `MaybeUninit::array_assume_init()`
        // once stable).
        Some(unsafe { std::mem::transmute_copy::<[MaybeUninit<T::Item>; N], [T::Item; N]>(&result) })
    }
}

(Note that this does not drop items from previous iterators in case next() panics. It is not required, but a good thing to do).

cafce25 · Accepted Answer · 2022-12-12 18:18:11Z

1

On nightly you can use this code. For stable you have to specifify the length in <[_; 4]>.

#![feature(generic_arg_infer)]
fn main() {
    let mut arr = [0u8..9, 2..8, 4..8, 3..5];
    for [a, b, c, d] in std::iter::from_fn(|| {
        <[_;_]>::try_from(
            arr.iter_mut()
            .map(|x| x.next())
            .collect::<Option<Vec<_>>>()?
            )
            .ok()
    }) {
        println!("{a} {b} {c} {d}");
    }
    println!("end of iteration");
    // 0 2 4 3
    // 1 3 5 4
    // end of iteration
}

answered Dec 12, 2022 at 18:18

cafce25

29.9k5 gold badges49 silver badges68 bronze badges

3 Comments

Chayim Friedman Over a year ago

Note that this collects into a Vec, so not so performant.

leo848 Over a year ago

@ChayimFriedman true, but until your method is stable I'd rather lose a bit of performance than introduce unsafe into my codebase.

Chayim Friedman Over a year ago

@leo848 That's definitely a right consideration, at least in some circumstances.

score 1 · Accepted Answer · 2022-12-13 22:18:00Z

1

On stable you can use this:

pub struct ZipArray<T, const N: usize> {
    array: [T; N],
}

pub fn zip_array<T: Iterator, const N: usize>(array: [T; N]) -> ZipArray<T, N> {
    ZipArray { array }
}

impl<T: Iterator, const N: usize> Iterator for ZipArray<T, N>
where
    <T as Iterator>::Item: Copy + Default,
{
    type Item = [T::Item; N];

    fn next(&mut self) -> Option<Self::Item> {
        let mut item = [T::Item::default(); N];

        for (index, element) in self.array.iter_mut().enumerate() {
            match element.next() {
                Some(value) => item[index] = value,
                None => {
                    return None;
                }
            }
        }

        Some(item)
    }
}

fn main() {
    let arr = [0..9, 2..8, 4..8, 3..5];
    for [a, b, c, d] in zip_array(arr) {
        println!("{a} {b} {c} {d}");
    }
    println!("end of iteration");
}

It's probably not as performant due to copies as the unsafe version and has additional trait bounds, but it's safe.

edited Dec 13, 2022 at 22:18

answered Dec 13, 2022 at 22:11

user3515670

1 Comment

leo848 Over a year ago

As far as I can tell, the trait bounds can be eliminated like this: Rust Playground - however I am not sure about the performance implications.

Collectives™ on Stack Overflow

How to turn an array of iterators into an iterator of arrays?

3 Answers 3

Comments

3 Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

3 Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related