`successors` is 2x slower than `map`

[Centri3]

I tried this code:

use std::iter::successors;

pub fn foo_map(y: u32) -> Vec<u32> {
    (0..).map(|i| y.pow(i)).take(5).collect()
}

pub fn foo_successors(y: u32) -> Vec<u32> {
    successors(Some(1), |x| Some(y * x)).take(5).collect()
}

godbolt

I expected to see this happen: the same performance

Instead, this happened: foo_successors is ~2x slower than foo_map, and the assembly outputted has a ton of checks. Changing foo_map to (0..).map(|i| Some(y.pow(i)).take(5).collect::<Option<_>>().unwrap() produces identical assembly (afaict, I have bad eyes :D), which leads me to believe it's because of Some, which should be optimized out as it never returns None.

Meta

rustc --version --verbose:

rustc 1.72.0-nightly (fe7454bf4 2023-06-19)
binary: rustc
commit-hash: fe7454bf439c93cbe9ac8a8f7fcfacd5a40244c2
commit-date: 2023-06-19
host: x86_64-unknown-linux-gnu
release: 1.72.0-nightly
LLVM version: 16.0.5

@rustbot label +I-slow

[the8472]

This is expected, Take<Map<RangeFrom>> has a tight size_hint and is TrustedLen while Successors does not and cannot have that because, as far as the the type system is concerned, it could return a None at any point.
Knowing the length of the iterator enables additional optimizations.
LLVM might eventually see that it is never None, but the type-based optimizations are applied earlier, at the library level, which make's LLVM's job much easier.

You'll likely get less efficient behavior too if you replace the map with filter_map.

[the8472]

It would require a variant of successors or from_fn where the closure returns T instead of Option<T> to get the same effect as (0..).map()

[quaternic]

repeat_with is pretty much that, you just have to keep the cached value as part of the closure:

pub fn foo_repeat_with(y: u32) -> Vec<u32> {
    let mut x = 1;
    repeat_with(|| { let r = x; x *= y; r } ).take(5).collect()
}

At least for this simple example, the generated code is essentially the same as for foo_map, and it does indeed avoid the check in e.g. .next().unwrap(): godbolt

[the8472]

I don't think there's much to do then. Except maybe improving documentation to direct users to the right method.

repeat_with even already has an example covering this case.