How do I avoid allocations in Iterator::flat_map? - rust

Question

I have a Vec of integers and I want to create a new Vec which contains those integers and squares of those integers. I could do this imperatively:
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);
but I want to use iterators. I came up with this code:
let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
.flat_map(|&x| vec![x, x * x])
.collect();
println!("{:?}", new_v);
but it allocates an intermediate Vec in the flat_map function.
How to use flat_map without allocations?

Answer 1

As of Rust 1.53.0, this can be written with just the array literal:
let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
.flat_map(|&x| [x, x * x])
.collect();
Rust 1.53.0 implements IntoIterator for arrays, so the vec![] and workarounds in previous solutions are no longer needed. This works on all editions.

Answer 2

If your iterator is small and you don't want any external dependencies, a short iterator can be constructed from std::iter::once and std::iter::Iterator::chain. For example,
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)
This could be made into a macro, though be aware that using this for too many elements may cause the recursion limit to be reached. If you're making an iterator for more than a few dozen elements, it's probably not too bad to have an allocation. If you really need the slight increase in performance, nnnmmm's solution is probably better.
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

As of version 1.51.0, the struct core::array::IntoIter has been stabilized. You can use it like this:
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();
The documentation warns that this may be deprecated in the future when IntoIterator is implemented for arrays, but currently it's the easiest way to do this.

Answer 4

You can use an ArrayVec for this.
let v = vec![1, 2, 3];
let new_v: Vec<_> = v.iter()
.flat_map(|&x| ArrayVec::from([x, x * x]))
.collect();
Making arrays be by-value iterators, so that you wouldn't need ArrayVec has been discussed, see https://github.com/rust-lang/rust/issues/25725 and the linked PRs.