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

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?

CodePudding user response：

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) })
    }
}

CodePudding user response：

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
}