使用Rayon实现const泛型并行Zip迭代器的with_producer问题

我希望从满足 I: IndexedParallelIterator 约束的数组 [I; N] 出发，创建一个元素类型为 [I::Item; N] 的 IndexedParallelIterator，将该迭代器命名为 ParConstZip<I, N>。

按照实现思路，完成该结构需要三个步骤：

• 实现对应的非并行版本 ConstZip<I, N> 迭代器，要求该迭代器实现 ExactSizeIterator 与 DoubleEndedIterator trait。
• 实现 Producer trait，用于拆分输入并生成 ConstZip<I, N> 迭代器实例。
• 基于上述 Producer，为 ParConstZip<I, N> 实现 ParallelIterator 与 IndexedParallelIterator trait。

目前前两步开发已经完成，代码如下：

非并行迭代器 ConstZip<I, N> 实现 #

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: 执行到此处时`dst`已完全初始化
        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: 执行到此处时`dst`已完全初始化
        unsafe { Some(MaybeUninit::array_assume_init(dst)) }
    }
}

Producer 实现 #

当前编写的Producer实现逻辑如下，符合对应trait的设计要求：

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: 长度为`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),
        )
    }
}

实现 IndexedParallelIterator 时遇到阻碍：除 with_producer 方法外其余逻辑均为样板代码可直接完成，但始终无法正确实现 with_producer 方法，当前 ParConstZip 代码如下：

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

此前曾就同类需求得到过Rayon库作者的高层实现思路指导，但反复查阅Rayon plumbing模块说明文档后，仍未理清 with_producer 方法的正确实现逻辑，现寻求对应技术解答。

内容的提问来源于stack exchange，提问作者MSR