如何对由任意迭代器构成的Vec实现类转置式并行迭代？

问题背景 #

我有一个形如[[1,2,3],[4,5,6],[7,8,9]]的嵌套Vec，希望实现一个函数some_func，使得执行以下代码时，按顺序输出[1,4,7]、[2,5,8]、[3,6,9]：

for k in some_func([[1,2,3],[4,5,6],[7,8,9]]) {
    println!("{k:?}")
}

目前我能手动实现针对有限长度嵌套Vec的转置，但想把下面的函数改造成通用版本，让参数支持Vec<I>（其中I: IntoIterator<Item = T>），适配任意迭代器，而非仅局限于Vec：

fn parallel_iterate<T>(items: Vec<Vec<T>>) -> Vec<Vec<T>>
where
    T: Clone,
{
    let size = items.iter().map(|row| row.len()).min().unwrap_or(0);
    let mut result = Vec::with_capacity(size);
    for index in 0..size {
        result.push(Vec::new());
        for inner in items.iter() {
            result[index].push(inner[index].clone())
        }
    }
    return result;
}

解决方案 #

针对任意迭代器的转置核心难点在于：迭代器只能被消费一次，且无法提前获取长度（除非实现ExactSizeIterator）。以下是通用实现方案：

通用转置迭代器实现 #

这个版本返回一个迭代器而非直接返回Vec，符合Rust迭代器范式，且支持所有实现IntoIterator的类型：

use std::iter::Peekable;

/// 转置一组迭代器，返回一个迭代器，每次产出一组对应位置的元素
fn transpose_iter<I, T>(mut items: Vec<I>) -> impl Iterator<Item = Vec<T>>
where
    I: IntoIterator<Item = T>,
{
    // 将所有输入迭代器转换为Peekable，方便检查是否还有剩余元素
    let mut iterators: Vec<Peekable<<I as IntoIterator>::IntoIter>> = items
        .drain(..)
        .map(|inner| inner.into_iter().peekable())
        .collect();

    std::iter::from_fn(move || {
        // 检查是否还有任何迭代器有未消费的元素
        if iterators.iter_mut().all(|it| it.peek().is_none()) {
            return None;
        }

        // 收集当前位置的所有元素
        let mut group = Vec::with_capacity(iterators.len());
        for it in iterators.iter_mut() {
            if let Some(item) = it.next() {
                group.push(item);
            }
        }

        Some(group)
    })
}

// 测试示例
fn main() {
    // 测试嵌套Vec
    let data = vec![vec![1,2,3], vec![4,5,6], vec![7,8,9]];
    for k in transpose_iter(data) {
        println!("{k:?}");
    }

    // 测试非Vec迭代器（如Range）
    let iter_data = vec![0..3, 3..6, 6..9];
    for k in transpose_iter(iter_data) {
        println!("{k:?}");
    }
}

关键细节说明 #

1. Peekable迭代器：用来提前检查每个内部迭代器是否还有剩余元素，避免无效的next()调用。
2. from_fn构造迭代器：通过闭包动态生成每一组元素，直到所有内部迭代器都被耗尽。
3. 兼容性：支持所有实现IntoIterator的类型，包括Vec、Range、HashMap的键/值迭代器等。

适配原函数的行为（仅取最短迭代器长度） #

如果需要和原函数保持一致，只收集到最短迭代器对应的元素数量，可以修改闭包逻辑：

std::iter::from_fn(move || {
    // 只要有一个迭代器耗尽，就停止生成元素
    if iterators.iter_mut().any(|it| it.peek().is_none()) {
        return None;
    }

    let mut group = Vec::with_capacity(iterators.len());
    for it in iterators.iter_mut() {
        // 这里unwrap安全，因为已经通过peek确认有元素
        group.push(it.next().unwrap());
    }

    Some(group)
})

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