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)

Details

Source code:

#![feature(try_trait, test)]externcrate test;use std::ops::Try;use test::{Bencher, black_box};#[no_mangle]pubfnstd_compute(a:u64,b:u64) -> u64{StepByWithoutTryFold(StepBy{iter: a..b,step:6,first_take:true,}).map(|x| x ^ (x - 1)).sum()}#[no_mangle]pubfnpr_compute(a:u64,b:u64) -> u64{StepBy{iter: a..b,step:6,first_take:true,}.map(|x| x ^ (x - 1)).sum()}#[bench]fnstd_bench(bencher:&mutBencher){let a = black_box(1);let b = black_box(5000000);
bencher.iter(|| {black_box(std_compute(a, b));});}#[bench]fnpr_bench(bencher:&mutBencher){let a = black_box(1);let b = black_box(5000000);
bencher.iter(|| {black_box(pr_compute(a, b));});}structStepBy<I>{iter:I,step:usize,first_take:bool,}structStepByWithoutTryFold<I>(StepBy<I>);impl<I:Iterator>IteratorforStepBy<I>{typeItem = I::Item;#[inline]fnnext(&mutself) -> Option<Self::Item>{ifself.first_take{self.first_take = false;self.iter.next()}else{self.iter.nth(self.step)}}#[inline]fntry_fold<B,F,R>(&mutself,init:B,mutf:F) -> RwhereSelf:Sized,F:FnMut(B,Self::Item) -> R,R:Try<Ok=B>{letmut accum = init;ifself.first_take{self.first_take = false;ifletSome(x) = self.iter.next(){
accum = f(accum, x)?;}else{returnTry::from_ok(accum);}}whileletSome(x) = self.iter.nth(self.step){
accum = f(accum, x)?;}Try::from_ok(accum)}}impl<I:Iterator>IteratorforStepByWithoutTryFold<I>{typeItem = I::Item;#[inline]fnnext(&mutself) -> Option<Self::Item>{ifself.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)

Details

#![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?