Ironing out StepBy<Range>'s performance issues

[Emerentius]

The behaviour of <Range<_> as Iterator>::nth has a slight mismatch with StepBy (or Step depending on your viewpoint) as @scottmcm has found out, resulting in sub-optimal performance.
On every iteration, the range has to first step forwards n-1 times to get the next element and then advance again by 1.

I'm hoping we can improve step_by into a 100% zero-cost abstraction.

It seems like the performance issue is specific to StepBy<Range>. I'm thinking therefore that we could specialize Iterator for StepBy<Range<I>> such that it would use @scottmcm's suggested semantics. Like this:

impl<I> Iterator for StepBy<Range<I>>
where
    I: Step
{
    fn next(&mut self) -> Option<Self::Item> {
        self.first = false;
        if let Some(mut n) = self.start.add_usize(self.step+1) {
            if n < self.end {
                std::mem::swap(&mut self.start, &mut n);
                return Some(n);
            }
        }

        self.start = self.end.clone();
        None
    }
}

That also avoids the branch on a regular next(). I haven't looked at the other methods but that boolean in StepBy could possibly become superfluous. During construction of the StepBy adapter, the size in .step_by(size) is decremented and this specialization has to counter-add 1 every time but that should be optimized away if inlined.

If someone were to depend on side-effects in Step::add_usize (when the trait is stabilized), this pre-stepping would become weird. Same thing with a hypothetical next_and_skip_ahead().

@scottmcm what do you think of this?

[scottmcm]

Overall yes; someone should make a PR 🙂 Probably best to use specialization for now to hide the detail. (I suppose a doc-hidden unstable method on Iterator could work too, but that feels wrong.)

A few minor things:

that we could specialize Iterator for StepBy<Range<I>>

I thought that normally public traits didn't get specialized, instead delegating to an internal one and specializing that, but apparently Iterator for Fuse is specialized, so I'm misremembering.

but that boolean in StepBy could possibly become superfluous

It would still be needed for non-Ranges, of course, so I don't think it can go away in the specialization approach, since Fuse still has its boolean even when the specialization never uses it.

Like this:

It's hard to guess at exactly how it should be written without looking at how LLVM handles it. Something like that seems plausible, but I could imagine seemingly-unimportant things like the branch orderings affecting the loop passes. Maybe it would be possible to have a codegen test to make sure?

[Emerentius]

I've copied the std implementations onto godbolt and added the specialization. As expected, it improves the generated code a lot.
You're right, @scottmcm, with switched branch order it produces slightly different, seemingly better, code. With that and usize ranges it compiles to nothing, u8 still has extra code. That is with constant start, end and step size, I haven't checked the rest, yet.

code on godbolt

Adding this specializion slightly changes what we're expecting from the Step trait in that any implementor must be fine with having 1 more add_usize() call than strictly necessary (side effects). It's unstable now anyway but in the future it may matter.
I'm going to send a PR for the specialization.

Is there any way we can get rid of the TryFrom baggage?

[scottmcm]

I haven't looked at it in detail yet, but reminder that this isn't exactly what we want to reach:

// what we want to reach
pub fn manual_while() {
    let mut n = 0;
    while n < UPPER {
        test::black_box(n);
        n += STEP;
    }
}

Because that doesn't have overflow detection, and is thus an infinite loop for something like (0..255u8).step_by(2), as it'll get to 254 and wrap back to 0, which are both still below 255.

[Emerentius]

Good catch. That's also what caused the difference between manual and specialized for u8.

[Emerentius]

Looks like specialization is still too broken for this. It trips up type inference. Fixed by second trait.

Another issue I've found is that with RangeFrom I can't specialize it generically in a way that inherits overflow checks based on debug or release. Furthermore, the overflow and therefore panic in debug mode would happen one iteration too early. We may still be able to specialize RangeFrom for each type separately whilst preserving semantics.

Step::add_usize for signed integers is currently impeding optimizations.

[the8472]

For the record, the issue is still open due because the specialization PR (#51601) was reverted in #56049

[the8472]

#111850 solved this for unsigned integers. Is that sufficient?