Refactor iteration logic in the Flatten and FlatMap iterators

[timvermeulen]

The Flatten and FlatMap iterators both delegate to FlattenCompat:

struct FlattenCompat<I, U> {
    iter: Fuse<I>,
    frontiter: Option<U>,
    backiter: Option<U>,
}

Every individual iterator method that FlattenCompat implements needs to carefully manage this state, checking whether the frontiter and backiter are present, and storing the current iterator appropriately if iteration is aborted. This has led to methods such as next, advance_by, and try_fold all having similar code for managing the iterator's state.

I have extracted this common logic of iterating the inner iterators with the option to exit early into a iter_try_fold method:

impl<I, U> FlattenCompat<I, U>
where
    I: Iterator<Item: IntoIterator<IntoIter = U>>,
{
    fn iter_try_fold<Acc, Fold, R>(&mut self, acc: Acc, fold: Fold) -> R
    where
        Fold: FnMut(Acc, &mut U) -> R,
        R: Try<Output = Acc>,
    { ... }
}

It passes each of the inner iterators to the given function as long as it keep succeeding. It takes care of managing FlattenCompat's state, so that the actual Iterator methods don't need to. The resulting code that makes use of this abstraction is much more straightforward:

fn next(&mut self) -> Option<U::Item> {
    #[inline]
    fn next<U: Iterator>((): (), iter: &mut U) -> ControlFlow<U::Item> {
        match iter.next() {
            None => ControlFlow::CONTINUE,
            Some(x) => ControlFlow::Break(x),
        }
    }

    self.iter_try_fold((), next).break_value()
}

Note that despite being implemented in terms of iter_try_fold, next is still able to benefit from U's next method. It therefore does not take the performance hit that implementing next directly in terms of Self::try_fold causes (in some benchmarks).

This PR also adds iter_try_rfold which captures the shared logic of try_rfold and advance_back_by, as well as iter_fold and iter_rfold for folding without early exits (used by fold, rfold, count, and last).

Benchmark results:

                                             before                after
bench_flat_map_sum                       423,255 ns/iter      414,338 ns/iter
bench_flat_map_ref_sum                 1,942,139 ns/iter    2,216,643 ns/iter
bench_flat_map_chain_sum               1,616,840 ns/iter    1,246,445 ns/iter
bench_flat_map_chain_ref_sum           4,348,110 ns/iter    3,574,775 ns/iter
bench_flat_map_chain_option_sum          780,037 ns/iter      780,679 ns/iter
bench_flat_map_chain_option_ref_sum    2,056,458 ns/iter      834,932 ns/iter

I added the last two benchmarks specifically to demonstrate an extreme case where FlatMap::next can benefit from custom internal iteration of the outer iterator, so take it with a grain of salt. We should probably do a perf run to see if the changes to next are worth it in practice.

[rust-highfive]

r? @m-ou-se

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

[rustbot]

Hey! It looks like you've submitted a new PR for the library teams!

If this PR contains changes to any rust-lang/rust public library APIs then please comment with @rustbot label +T-libs-api -T-libs to tag it appropriately. If this PR contains changes to any unstable APIs please edit the PR description to add a link to the relevant API Change Proposal or create one if you haven't already. If you're unsure where your change falls no worries, just leave it as is and the reviewer will take a look and make a decision to forward on if necessary.

Examples of T-libs-api changes:

• Stabilizing library features
• Introducing insta-stable changes such as new implementations of existing stable traits on existing stable types
• Introducing new or changing existing unstable library APIs (excluding permanently unstable features / features without a tracking issue)
• Changing public documentation in ways that create new stability guarantees
• Changing observable runtime behavior of library APIs

[the8472]

r? @the8472

[the8472]

The general DRYing seems fine, but do the last() and count() impls provide any benefit? The default impls already use fold.

bench_flat_map_ref_sum 1,942,139 ns/iter 2,216,643 ns/iter

This seems to be a slight regression. Since by-ref iterators don't forward internal iteration methods and rely on next() instead does that mean next() is now a bit less efficient?
Implementing next by using folding over an inner iterator is unusual and potentially means more codegen as the try_fold on adapters have a deeper callgraph compared to next implementations.
Maybe it'd help in cases where the outer iterator has to walk over tons of empty inner elements before it arrives at a non-empty one, but seems like an uncommon scenario when stepping with next() [Edit: I see that's what bench_flat_map_chain_option_ref_sum demonstrates]

Anyway, let's do a perf run since the compiler does use flatten and flat_map quite a bit

[the8472]

@bors try @rust-timer queue

[rust-timer]

Awaiting bors try build completion.

@rustbot label: +S-waiting-on-perf

[bors]

⌛ Trying commit f269998e5d0b44dd9fd78d4573c4386f7d07ad93 with merge 06aa8c36743913f0893953065baf3bb61c5f0d29...

[bors]

☀️ Try build successful - checks-actions
Build commit: 06aa8c36743913f0893953065baf3bb61c5f0d29 (06aa8c36743913f0893953065baf3bb61c5f0d29)

[rust-timer]

Queued 06aa8c36743913f0893953065baf3bb61c5f0d29 with parent 47ba935, future comparison URL.

[timvermeulen]

The general DRYing seems fine, but do the last() and count() impls provide any benefit? The default impls already use fold.

The default impls use Self::fold, which in turn uses I::fold and U::fold. This custom last can benefit from a custom last implementation on U. Calling flatten().last() on something like a Once<Chunks<'_, T>> demonstrates this well, since the optimizer seems unable to figure out what the last chunk is based on the default implementation of Iterator::last alone.

bench_flat_map_ref_sum 1,942,139 ns/iter 2,216,643 ns/iter

This seems to be a slight regression. Since by-ref iterators don't forward internal iteration methods and rely on next() instead does that mean next() is now a bit less efficient?

In this particular benchmark, yes. Not in the other two by-ref benchmarks. It's entirely possible that this one is more representative of real world code though.

Implementing next by using folding over an inner iterator is unusual and potentially means more codegen as the try_fold on adapters have a deeper callgraph compared to next implementations.

This is primarily true when only a single call to the inner next is needed. An iterator like FilterMap does (indirectly) call try_fold on the inner iterator from its next implementation. You're right about this only being beneficial when FlattenCompat::next has to skip over empty Us and therefore advance I by more than one.

[rust-timer]

Finished benchmarking commit (06aa8c36743913f0893953065baf3bb61c5f0d29): comparison url.

Instruction count

• Primary benchmarks: mixed results
• Secondary benchmarks: 🎉 relevant improvements found

	mean1	max	count2
Regressions 😿 (primary)	1.0%	1.9%	11
Regressions 😿 (secondary)	N/A	N/A	0
Improvements 🎉 (primary)	-0.5%	-0.7%	8
Improvements 🎉 (secondary)	-0.5%	-1.0%	8
All 😿🎉 (primary)	0.4%	1.9%	19

Max RSS (memory usage)

Results

• Primary benchmarks: 🎉 relevant improvements found
• Secondary benchmarks: mixed results

	mean1	max	count2
Regressions 😿 (primary)	N/A	N/A	0
Regressions 😿 (secondary)	2.6%	3.1%	5
Improvements 🎉 (primary)	-4.8%	-7.0%	3
Improvements 🎉 (secondary)	-2.0%	-2.0%	1
All 😿🎉 (primary)	-4.8%	-7.0%	3

Cycles

Results

• Primary benchmarks: mixed results
• Secondary benchmarks: no relevant changes found

	mean1	max	count2
Regressions 😿 (primary)	2.6%	2.8%	2
Regressions 😿 (secondary)	N/A	N/A	0
Improvements 🎉 (primary)	-2.4%	-2.4%	1
Improvements 🎉 (secondary)	N/A	N/A	0
All 😿🎉 (primary)	0.9%	2.8%	3

If you disagree with this performance assessment, please file an issue in rust-lang/rustc-perf.

Benchmarking this pull request likely means that it is perf-sensitive, so we're automatically marking it as not fit for rolling up. While you can manually mark this PR as fit for rollup, we strongly recommend not doing so since this PR may lead to changes in compiler perf.

Next Steps: If you can justify the regressions found in this try perf run, please indicate this with @rustbot label: +perf-regression-triaged along with sufficient written justification. If you cannot justify the regressions please fix the regressions and do another perf run. If the next run shows neutral or positive results, the label will be automatically removed.

@bors rollup=never
@rustbot label: +S-waiting-on-review -S-waiting-on-perf +perf-regression

Footnotes

1. the arithmetic mean of the percent change ↩ ↩² ↩³

2. number of relevant changes ↩ ↩² ↩³

[the8472]

So the perf numbers seem more negative than positive and it spends more time in codegen (even in debug) and in llvm opts. And the artifact sizes are larger too. I think it'd make sense to look at the llvm-lines output or the generated assembly.

[timvermeulen]

I think there's not much we can do here to improve this. If these perf results aren't deemed worth it for speeding up external iteration of flattened iterators where many of them are empty, then I think we should revert the changes to FlattenCompat::next.

[the8472]

Sure, we can measure just the fold()-changes without the next() ones.

If these perf results aren't deemed worth it for speeding up external iteration of flattened iterators where many of them are empty

The microbenchmarks look worse for dense outer iterators, the perf results look worse on several metrics so I think we'd at least need some evidence that these kinds of iterators are common and that the speedups are worth it. And I think that is exactly the kind of scenario where advance_by(0) as a kind of optimization hint would be useful.

[timvermeulen]

@bors try @rust-timer queue

[rust-timer]

Insufficient permissions to issue commands to rust-timer.

[bors]

@timvermeulen: 🔑 Insufficient privileges: not in try users

[the8472]

@bors try @rust-timer queue

[rust-timer]

Awaiting bors try build completion.

@rustbot label: +S-waiting-on-perf

[bors]

⌛ Trying commit cce733e6886c7bbac2e6d41add17a6645d4cf189 with merge ca6848a65861be4083fc65253d34ecf5f794338f...

[timvermeulen]

And I think that is exactly the kind of scenario where advance_by(0) as a kind of optimization hint would be useful.

It makes more sense for Skip/Take than for Flatten. A flattened iterator can have empty iterators throughout, not just at the start/end. For this reason, this kind of optimization makes more sense inside next (if it were worthwhile).

[bors]

☀️ Try build successful - checks-actions
Build commit: ca6848a65861be4083fc65253d34ecf5f794338f (ca6848a65861be4083fc65253d34ecf5f794338f)

[rust-timer]

Queued ca6848a65861be4083fc65253d34ecf5f794338f with parent c9e134e, future comparison URL.

[rust-timer]

Finished benchmarking commit (ca6848a65861be4083fc65253d34ecf5f794338f): comparison url.

Instruction count

This benchmark run did not return any relevant results for this metric.

Max RSS (memory usage)

Results

• Primary benchmarks: 🎉 relevant improvements found
• Secondary benchmarks: 🎉 relevant improvement found

	mean1	max	count2
Regressions 😿 (primary)	N/A	N/A	0
Regressions 😿 (secondary)	N/A	N/A	0
Improvements 🎉 (primary)	-2.2%	-2.7%	3
Improvements 🎉 (secondary)	-2.0%	-2.0%	1
All 😿🎉 (primary)	-2.2%	-2.7%	3

Cycles

Results

• Primary benchmarks: no relevant changes found
• Secondary benchmarks: 😿 relevant regression found

	mean1	max	count2
Regressions 😿 (primary)	N/A	N/A	0
Regressions 😿 (secondary)	2.4%	2.4%	1
Improvements 🎉 (primary)	N/A	N/A	0
Improvements 🎉 (secondary)	N/A	N/A	0
All 😿🎉 (primary)	N/A	N/A	0

If you disagree with this performance assessment, please file an issue in rust-lang/rustc-perf.

Benchmarking this pull request likely means that it is perf-sensitive, so we're automatically marking it as not fit for rolling up. While you can manually mark this PR as fit for rollup, we strongly recommend not doing so since this PR may lead to changes in compiler perf.

@bors rollup=never
@rustbot label: +S-waiting-on-review -S-waiting-on-perf -perf-regression

Footnotes

1. the arithmetic mean of the percent change ↩ ↩²

2. number of relevant changes ↩ ↩²

[timvermeulen]

I'm happy to have this land as is.

[the8472]

@bors r+ rollup=never

[bors]

📌 Commit 38bb0b1 has been approved by the8472

It is now in the queue for this repository.

[bors]

⌛ Testing commit 38bb0b1 with merge 6c943ba...

[bors]

☀️ Test successful - checks-actions
Approved by: the8472
Pushing 6c943ba to master...

[rust-timer]

Finished benchmarking commit (6c943ba): comparison url.

Instruction count

This benchmark run did not return any relevant results for this metric.

Max RSS (memory usage)

Results

• Primary benchmarks: ❌ relevant regression found
• Secondary benchmarks: mixed results

	mean1	max	count2
Regressions ❌ (primary)	2.1%	2.1%	1
Regressions ❌ (secondary)	1.6%	1.8%	4
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-2.4%	-2.9%	3
All ❌✅ (primary)	2.1%	2.1%	1

Cycles

This benchmark run did not return any relevant results for this metric.

If you disagree with this performance assessment, please file an issue in rust-lang/rustc-perf.

@rustbot label: -perf-regression

Footnotes

1. the arithmetic mean of the percent change ↩

2. number of relevant changes ↩