implement fold() on array::IntoIter to improve flatten().collect() perf

[the8472]

With #87168 flattening array::IntoIters is now TrustedLen, the FromIterator implementation for Vec has a specialization for TrustedLen iterators which uses internal iteration. This implements one of the main internal iteration methods on array::Into to optimize the combination of those two features.

This should address the main issue in #87411

# old
test vec::bench_flat_map_collect                         ... bench:   2,244,024 ns/iter (+/- 18,903)

# new
test vec::bench_flat_map_collect                         ... bench:     172,863 ns/iter (+/- 2,141)

[rust-highfive]

r? @kennytm

(rust-highfive has picked a reviewer for you, use r? to override)

[kennytm]

while this LGTM, shouldn't the original issue be addressed by implementing SpecExtend<T, std::array::IntoIter<T>> for Vec<T, A>?

[the8472]

The original issue involved Flatten which results in several adapters sitting between SpecExtend and the IntoIter.

[kennytm]

can we call fold here instead of try_fold?

Suggested change

	(&mut self.alive)
	.try_fold::<_, _, Result<_, !>>(init, |acc, idx| {
	// SAFETY: idx is obtained by folding over the `alive` range, which implies the
	// value is currently considered alive but as the range is being consumed each value
	// we read here will only be read once and then considered dead.
	Ok(fold(acc, unsafe { data.get_unchecked(idx).assume_init_read() }))
	})
	.unwrap()
	self.alive.fold(init, |acc, idx| {
	// SAFETY: idx is obtained by folding over the `alive` range, which implies the
	// value is currently considered alive but as the range is being consumed each value
	// we read here will only be read once and then considered dead.
	fold(acc, unsafe { data.get_unchecked(idx).assume_init_read() })
	})

[the8472]

I tried that, array::IntoIter has a Drop impl, so alive can't be move out, but that would be required to call fold(self), that's why I used try_fold instead.

[kennytm]

@the8472 oops right.

(&mut self.alive).fold(init, ...) should work though.

[the8472]

Yeah but that would go through impl Iterator for &mut I which is less optimized.

[kennytm]

self.alive is a std::ops::Range<usize> and AFAIK there is no special-cased implementation of fold or try_fold for Range<usize> nor &mut Range<usize>.

[kennytm]

even if we do the mem::take it will just turn self.alive to 0..0 and then leaks everything which is safe 🙃 (compared with self.alive.clone().fold(...) which will cause double-free).

[the8472]

I have now tried (&mut self.alive).fold instead of try_fold, it undoes all perfomance gains. I guess somehow the indirection through &mut inhibits optimizations.

[the8472]

Which is due to this not being #[inline]

rust/library/core/src/iter/traits/iterator.rs

Lines 3474 to 3478 in 71a6c7c

	impl<I: Iterator + ?Sized> Iterator for &mut I {
	type Item = I::Item;
	fn next(&mut self) -> Option<I::Item> {
	(**self).next()
	}

[kennytm]

heh.

wdyt should we just add the #[inline] or leave a FIXME comment explaining the performance regression if we use fold instead of try_fold? either way is fine for me.

[the8472]

I'll add a FIXME, changing inlining on such central methods can have mixed impact on compile time even if runtime performance is better, so that should be done on a separate PR.

[kennytm]

@bors r+ rollup=iffy

[bors]

📌 Commit 2276c5e has been approved by kennytm

[bors]

⌛ Testing commit 2276c5e with merge 99d6692...

[bors]

☀️ Test successful - checks-actions
Approved by: kennytm
Pushing 99d6692 to master...