Writing a const generic parallel zip using rayon - rust

Question

I wish to create an IndexedParallelIterator<Item = [I::Item; N]> from an array [I; N] where I: IndexedParallelIterator. Let's call it ParConstZip<I, N>.
As I understand, there's a couple of steps here:
Create the corresponding non-parallel ConstZip<I, N> iterator (which must implement ExactSizeIterator and DoubleEndedIterator).
Create a Producer to split the input and create the ConstZip<I, N> iterators.
Implement ParallelIterator and IndexedParallelIterator for ParConstZip<I, N> using the Producer.
I think I'm alright on points 1. and 2. Specifically, here's the ConstZip<I, N> implementation:
pub struct ConstZip<I, const N: usize>([I; N]);
impl<I, const N: usize> Iterator for ConstZip<I, N>
where
I: Iterator,
{
type Item = [I::Item; N];
fn next(&mut self) -> Option<Self::Item> {
let mut dst = MaybeUninit::uninit_array();
for (i, iter) in self.0.iter_mut().enumerate() {
dst[i] = MaybeUninit::new(iter.next()?);
}
// SAFETY: If we reach this point, `dst` has been fully initialized
unsafe { Some(MaybeUninit::array_assume_init(dst)) }
}
}
impl<I, const N: usize> ExactSizeIterator for ConstZip<I, N>
where
I: ExactSizeIterator,
{
fn len(&self) -> usize {
self.0.iter().map(|x| x.len()).min().unwrap()
}
}
impl<I, const N: usize> DoubleEndedIterator for ConstZip<I, N>
where
I: DoubleEndedIterator,
{
fn next_back(&mut self) -> Option<Self::Item> {
let mut dst = MaybeUninit::uninit_array();
for (i, iter) in self.0.iter_mut().enumerate() {
dst[i] = MaybeUninit::new(iter.next_back()?);
}
// SAFETY: If we reach this point, `dst` has been fully initialized
unsafe { Some(MaybeUninit::array_assume_init(dst)) }
}
}
And here's what I believe to be an appropriate producer:
pub struct ParConstZipProducer<P, const N: usize>([P; N]);
impl<P, const N: usize> Producer for ParConstZipProducer<P, N>
where
P: Producer,
{
type Item = [P::Item; N];
type IntoIter = ConstZip<P::IntoIter, N>;
fn into_iter(self) -> Self::IntoIter {
ConstZip(self.0.map(Producer::into_iter))
}
fn split_at(self, index: usize) -> (Self, Self) {
let mut left_array = MaybeUninit::uninit_array();
let mut right_array = MaybeUninit::uninit_array();
for (i, producer) in self.0.into_iter().enumerate() {
let (left, right) = producer.split_at(index);
left_array[i] = MaybeUninit::new(left);
right_array[i] = MaybeUninit::new(right);
}
// SAFETY: Arrays are guaranteed to be fully initialised at length `N`
let left_array = unsafe { MaybeUninit::array_assume_init(left_array) };
let right_array = unsafe { MaybeUninit::array_assume_init(right_array) };
(
ParConstZipProducer(left_array),
ParConstZipProducer(right_array),
)
}
}
However, I stumble when it comes to the actual implementation of IndexedParallelIterator. Most of it seems to be boilerplate, but the with_producer method I cannot figure out how to implement:
pub struct ParConstZip<I, const N: usize>([I; N]);
impl<I, const N: usize> ParallelIterator for ParConstZip<I, N>
where
I: IndexedParallelIterator,
{
type Item = [I::Item; N];
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
bridge(self, consumer)
}
}
impl<I, const N: usize> IndexedParallelIterator for ParConstZip<I, N>
where
I: IndexedParallelIterator,
{
fn drive<C>(self, consumer: C) -> C::Result
where
C: Consumer<Self::Item>,
{
bridge(self, consumer)
}
fn len(&self) -> usize {
self.0.iter().map(|x| x.len()).min().unwrap()
}
fn with_producer<CB>(self, callback: CB) -> CB::Output
where
CB: ProducerCallback<Self::Item>,
{
todo!()
}
}
Playground link.
I needed the same thing about a year ago, and asked a question on the users.rust-lang forum. I got some helpful high-level pointers from one of the authors of rayon, but after reading the rayon plumbing README a number of times both then and now, I have to admit I'm still a bit lost.