Rust：避免在“ 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);

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

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

Answer 1

您可以为此使用ArrayVec。

ArrayVec

使数组成为按值迭代器，因此您无需讨论

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

，请参阅ArrayVec和链接的PR。

Answer 2

如果您的迭代器很小，并且您不希望任何外部依赖项，则可以根据https://github.com/rust-lang/rust/issues/25725和std::iter::once构造一个简短的迭代器。例如，

std::iter::once

std::iter::Iterator::chain

这可以做成一个宏，尽管要知道，对太多元素使用它可能会导致达到递归限制。如果要为多个元素创建一个迭代器，则分配可能并不算太糟。如果您确实需要稍微提高性能，则nnnmmm的解决方案可能更好。

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);