CODEWITHSUNDEEP
arraysrustiterator
leo848
leo848

Reputation: 667

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

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?

Upvotes: 0

Views: 435

Answers (3)

Richard Neumann
Richard Neumann

Reputation: 3361

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.

Upvotes: 1

Chayim Friedman
Chayim Friedman

Reputation: 70970

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

Upvotes: 2

cafce25
cafce25

Reputation: 27357

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
}

Upvotes: 1

Related Questions