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?