如何避免在 Iterator::flat_map 中进行分配？

Question

我有一个 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？

Answer 1

您可以为此使用 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。

Answer 2

如果您的迭代器很小并且您不需要任何外部依赖项，可以从 std::iter::once and 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)

Answer 3

从版本 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 时，这可能会在未来被弃用，但目前这是最简单的方法。

Answer 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![] 和解决方法。这适用于所有版本。

如何避免在 Iterator::flat_map 中进行分配？

How do I avoid allocations in Iterator::flat_map?

rust

iterator

declarative