Rust中能否将同一结构的多个drain迭代器展平？有何可行实现方案？

算力阿峰

2026-9-27

错误本质

Rust借用检查器无法证明flat_map回调中返回的drain迭代器不会与后续回调对vals的可变访问产生冲突，即使逻辑上flat_map会消费完前一个迭代器才会调用下一次回调，静态检查也无法识别这种运行时保证。

可行方案

方案1：预收集结果（无unsafe，推荐大多数场景使用）

先将所有需要取出的元素收集到临时容器中，再返回迭代器，同时需要先对索引做去重倒序处理，避免前面的元素移除导致后续索引错位：

fn drain_some<'a>(
    vals: &'a mut Vec<Vec<i32>>,
    inds: &'a [usize],
) -> impl Iterator<Item = i32> + 'a {
    // 索引去重+倒序排序，避免删除前序元素导致后续索引失效
    let mut sorted_inds = inds.to_vec();
    sorted_inds.sort_unstable_by(|a, b| b.cmp(a));
    sorted_inds.dedup();

    let mut result = Vec::new();
    for idx in sorted_inds {
        result.extend(vals[idx].drain(..));
    }
    result.into_iter()
}

优点：完全符合安全 Rust 规则，无额外心智负担
缺点：需要额外分配内存存储结果，大数据量场景下有一定开销

方案2：Unsafe封装（零额外开销，适合性能敏感场景）

如果可以保证索引唯一、且不会重复访问同一块内存，可以通过UnsafeCell封装绕过借用检查，自己维护借用安全：

use std::cell::UnsafeCell;

fn drain_some<'a>(
    vals: &'a mut Vec<Vec<i32>>,
    inds: &'a [usize],
) -> impl Iterator<Item = i32> + 'a {
    // 安全保证：我们只会同时持有一个子Vec的可变引用，且索引无重复
    let vals = UnsafeCell::new(vals.as_mut_slice());
    inds.iter().flat_map(move |&i| {
        unsafe { &mut *vals.get() }[i].drain(..)
    })
}