如何避免 Iterator::flat_map 中的分配？答案

【问题标题】：How do I avoid allocations in Iterator::flat_map?如何避免 Iterator::flat_map 中的分配？
【发布时间】：2019-11-30 10:35:46
【问题描述】：

我有一个整数Vec，我想创建一个新的Vec，其中包含这些整数和这些整数的平方。我可以强制执行：

let v = vec![1, 2, 3];
let mut new_v = Vec::new(); // new instead of with_capacity for simplicity sake.
for &x in v.iter() {
    new_v.push(x);
    new_v.push(x * x);
}
println!("{:?}", new_v);

但我想使用迭代器。我想出了这段代码：

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| vec![x, x * x])
    .collect();
println!("{:?}", new_v);

但它在flat_map 函数中分配了一个中间Vec。

如何在没有分配的情况下使用flat_map？

【问题讨论】：

标签： rust iterator declarative

【解决方案1】：

从 1.51.0 版本开始，结构 core::array::IntoIter 已经稳定。你可以这样使用它：

use core::array;

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| array::IntoIter::new([x, x * x]))
    .collect();

文档警告说，将来当为数组实现 IntoIterator 时，这可能会被弃用，但目前这是最简单的方法。

【讨论】：

需要明确的是，可能被弃用的“this”是使用new 函数创建的，而不是array::IntoIter 本身。如果数组确实得到IntoIterator，你就可以写flat_map(|&x| [x, x * x])。

【解决方案2】：

您可以为此使用ArrayVec。

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| ArrayVec::from([x, x * x]))
    .collect();

使数组成为按值迭代器，这样您就不需要 ArrayVec 已经讨论过了，请参阅 https://github.com/rust-lang/rust/issues/25725 和链接的 PR。

【讨论】：

【解决方案3】：

如果您的迭代器很小并且您不想要任何外部依赖项，则可以从std::iter::once 和std::iter::Iterator::chain 构造一个短迭代器。例如，

use std::iter;

let v = vec![1, 2, 3];
let new_v: Vec<_> = v
    .iter()
    .flat_map(|&x| iter::once(x).chain(iter::once(x * x)))
    .collect();
println!("{:?}", new_v);

(playground)

这可以做成一个宏，但要注意对太多元素使用它可能会导致达到递归限制。如果您要为超过几十个元素创建迭代器，那么进行分配可能还不错。如果你真的需要性能的轻微提升，nnnmmm 的解决方案可能会更好。

macro_rules! small_iter {
    () => { std::iter::empty() };
    ($x: expr) => {
        std::iter::once($x)
    };
    ($x: expr, $($y: tt)*) => {
        std::iter::once($x).chain(small_iter!($($y)*))
    };
}

fn main() {
    let v = vec![1, 2, 3];
    let new_v: Vec<_> = v
        .iter()
        .flat_map(|&x| small_iter!(x, x * x))
        .collect();
    println!("{:?}", new_v);
}

(playground)

【讨论】：

【解决方案4】：

从 Rust 1.53.0 开始，这可以只使用数组字面量来编写：

let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
    .flat_map(|&x| [x, x * x])
    .collect();

Rust 1.53.0 实现了IntoIterator for arrays，因此不再需要以前解决方案中的vec![] 和解决方法。这适用于所有版本。

【讨论】：