Override try_fold for StepBy

[SimonSapin]

No description provided.

[rust-highfive]

r? @kennytm

(rust_highfive has picked a reviewer for you, use r? to override)

[SimonSapin]

@scottmcm This PR is from your repo. I’m not good at reading assembly but your comment at #27741 (comment) suggests that this is an improvement. Should we land it?

The diff looks good to me.

[kennytm]

I've done a simple benchmark and the test code from the comment is improved from 800µs/iter to 600µs/iter.

running 2 tests
test pr_bench  ... bench:     592,516 ns/iter (+/- 42,563)
test std_bench ... bench:     809,164 ns/iter (+/- 54,219)

Source code:

#![feature(try_trait, test)]

extern crate test;

use std::ops::Try;

use test::{Bencher, black_box};

#[no_mangle]
pub fn std_compute(a: u64, b: u64) -> u64 {
    StepByWithoutTryFold(StepBy {
        iter: a..b,
        step: 6,
        first_take: true,
    }).map(|x| x ^ (x - 1)).sum()
}

#[no_mangle]
pub fn pr_compute(a: u64, b: u64) -> u64 {
    StepBy {
        iter: a..b,
        step: 6,
        first_take: true,
    }.map(|x| x ^ (x - 1)).sum()
}

#[bench]
fn std_bench(bencher: &mut Bencher) {
    let a = black_box(1);
    let b = black_box(5000000);
    bencher.iter(|| {
        black_box(std_compute(a, b));
    });
}

#[bench]
fn pr_bench(bencher: &mut Bencher) {
    let a = black_box(1);
    let b = black_box(5000000);
    bencher.iter(|| {
        black_box(pr_compute(a, b));
    });
}

struct StepBy<I> {
    iter: I,
    step: usize,
    first_take: bool,
}

struct StepByWithoutTryFold<I>(StepBy<I>);

impl<I: Iterator> Iterator for StepBy<I> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.first_take {
            self.first_take = false;
            self.iter.next()
        } else {
            self.iter.nth(self.step)
        }
    }

    #[inline]
    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R where
        Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
    {
        let mut accum = init;

        if self.first_take {
            self.first_take = false;
            if let Some(x) = self.iter.next() {
                accum = f(accum, x)?;
            } else {
                return Try::from_ok(accum);
            }
        }

        while let Some(x) = self.iter.nth(self.step) {
            accum = f(accum, x)?;
        }
        Try::from_ok(accum)
    }
}

impl<I: Iterator> Iterator for StepByWithoutTryFold<I> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.0.first_take {
            self.0.first_take = false;
            self.0.iter.next()
        } else {
            self.0.iter.nth(self.0.step)
        }
    }
}

[scottmcm]

Hmm, I'd forgotten I'd started this 😅 I see two options:

1. Just merge this, since it helps -- thanks for benching, @kennytm! (Well, with some tests, especially for the short-circuiting logic, since that's the thing I got wrong the most doing the other try_folds.)

2. Reimplement StepBy so that the bool flag isn't needed, and thus the manual try_fold is likely also unneeded (since there's nothing to manually hoist). The flag is there so it can use nth, but if there was a method that called next n times, returning the first result, then the flag wouldn't be needed, and it would actually be easier for Range to implement. Sketch of what I mean: https://play.rust-lang.org/?gist=f2c028e6e42e9ce5f3a311b71dcfc3a9&version=nightly&mode=release

Thoughts?

[kennytm]

I think we could do the refactoring later.

@bors r+

[bors]

📌 Commit 2d55d28 has been approved by kennytm

[bors]

🔒 Merge conflict

[Emerentius]

I've rerun @kennytm's code but replaced the std code by the specialization from my PR referencing this issue just above this comment. With that, this PR would lower performance.

test pr_bench             ... bench:     647,483 ns/iter (+/- 12,094)
test specialized_pr_bench ... bench:     413,365 ns/iter (+/- 13,795)

#![feature(try_trait, test, step_trait)]
use std::iter::Step;

extern crate test;

use std::ops::Try;

use test::{Bencher, black_box};

#[no_mangle]
pub fn std_compute(a: u64, b: u64) -> u64 {
    StepByWithoutTryFold(StepBy {
        iter: a..b,
        step: 6,
        first_take: true,
    }).map(|x| x ^ (x - 1)).sum()
}

#[no_mangle]
pub fn pr_compute(a: u64, b: u64) -> u64 {
    StepBy {
        iter: a..b,
        step: 6,
        first_take: true,
    }.map(|x| x ^ (x - 1)).sum()
}

#[bench]
fn specialized_pr_bench(bencher: &mut Bencher) {
    let a = black_box(1);
    let b = black_box(5000000);
    bencher.iter(|| {
        black_box(std_compute(a, b));
    });
}

#[bench]
fn pr_bench(bencher: &mut Bencher) {
    let a = black_box(1);
    let b = black_box(5000000);
    bencher.iter(|| {
        black_box(pr_compute(a, b));
    });
}

struct StepBy<I> {
    iter: I,
    step: usize,
    first_take: bool,
}

struct StepByWithoutTryFold<I>(StepBy<I>);

impl<I: Iterator> Iterator for StepBy<I> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.first_take {
            self.first_take = false;
            self.iter.next()
        } else {
            self.iter.nth(self.step)
        }
    }

    #[inline]
    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R where
        Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
    {
        let mut accum = init;

        if self.first_take {
            self.first_take = false;
            if let Some(x) = self.iter.next() {
                accum = f(accum, x)?;
            } else {
                return Try::from_ok(accum);
            }
        }

        while let Some(x) = self.iter.nth(self.step) {
            accum = f(accum, x)?;
        }
        Try::from_ok(accum)
    }
}

impl<T> Iterator for StepByWithoutTryFold<std::ops::Range<T>>
where
    T: Step,
{
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.0.first_take = false;
        if self.0.iter.start >= self.0.iter.end {
            return None;
        }
        // add 1 to self.step to get original step size back
        // it was decremented for the general case on construction
        if let Some(n) = self.0.iter.start.add_usize(self.0.step+1) {
            let next = std::mem::replace(&mut self.0.iter.start, n);
            Some(next)
        } else {
            let last = self.0.iter.start.clone();
            self.0.iter.start = self.0.iter.end.clone();
            Some(last)
        }
    }
}

[SimonSapin]

Thanks for looking into this @Emerentius. Closing in favor of #51601.

[timvermeulen]

Since #51601 was reverted, should this be considered again?