diff --git a/src/libcore/benches/iter.rs b/src/libcore/benches/iter.rs index 6c597301ac204..36055d0ab01db 100644 --- a/src/libcore/benches/iter.rs +++ b/src/libcore/benches/iter.rs @@ -310,3 +310,39 @@ fn bench_skip_then_zip(b: &mut Bencher) { assert_eq!(s, 2009900); }); } + +#[bench] +fn bench_triples_inclusive_range(b: &mut Bencher) { + // Example from #45222 + b.iter(|| { + (1..) + .flat_map(|z| { + (1u32..=z).flat_map(move |x| { + (x..=z) + .filter(move |&y| x.pow(2) + y.pow(2) == z.pow(2)) + .map(move |y| (x, y, z)) + }) + }) + .take(20) + .map(|(x, y, z)| x + y + z) + .sum::() + }) +} + +#[bench] +fn bench_triples_range(b: &mut Bencher) { + // Example from #45222 + b.iter(|| { + (1..) + .flat_map(|z| { + (1u32..z+1).flat_map(move |x| { + (x..z+1) + .filter(move |&y| x.pow(2) + y.pow(2) == z.pow(2)) + .map(move |y| (x, y, z)) + }) + }) + .take(20) + .map(|(x, y, z)| x + y + z) + .sum::() + }) +} diff --git a/src/libcore/iter/range.rs b/src/libcore/iter/range.rs index f0fd07b43cae0..ff6723fa2666b 100644 --- a/src/libcore/iter/range.rs +++ b/src/libcore/iter/range.rs @@ -10,7 +10,7 @@ use convert::TryFrom; use mem; -use ops::{self, Add, Sub}; +use ops::{self, Add, Sub, Try}; use usize; use super::{FusedIterator, TrustedLen}; @@ -270,6 +270,23 @@ impl Iterator for ops::Range { fn max(mut self) -> Option { self.next_back() } + + #[inline] + fn try_fold(&mut self, init: B, mut f: F) -> R where + Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try + { + let mut accum = init; + if self.start >= self.end { + return Try::from_ok(accum); + } + + while self.start < self.end { + let n = self.start.add_one(); + accum = f(accum, mem::replace(&mut self.start, n))?; + } + + Try::from_ok(accum) + } } // These macros generate `ExactSizeIterator` impls for various range types. @@ -406,6 +423,26 @@ impl Iterator for ops::RangeInclusive { fn max(mut self) -> Option { self.next_back() } + + #[inline] + fn try_fold(&mut self, init: B, mut f: F) -> R where + Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try + { + let mut accum = init; + + self.compute_is_empty(); + if self.is_empty.unwrap_or_default() { + return Try::from_ok(accum); + } + + while self.start < self.end { + let n = self.start.add_one(); + accum = f(accum, mem::replace(&mut self.start, n))?; + } + + self.is_empty = Some(true); + f(accum, self.start.clone()) + } } #[stable(feature = "inclusive_range", since = "1.26.0")] @@ -425,6 +462,26 @@ impl DoubleEndedIterator for ops::RangeInclusive { self.end.clone() }) } + + #[inline] + fn try_rfold(&mut self, init: B, mut f: F) -> R where + Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try + { + let mut accum = init; + + self.compute_is_empty(); + if self.is_empty.unwrap_or_default() { + return Try::from_ok(accum); + } + + while self.start < self.end { + let n = self.end.sub_one(); + accum = f(accum, mem::replace(&mut self.end, n))?; + } + + self.is_empty = Some(true); + f(accum, self.end.clone()) + } } #[stable(feature = "fused", since = "1.26.0")]