From 1b1e887f4d40e5800a8d5ae81f8574806e7ba21a Mon Sep 17 00:00:00 2001 From: Scott McMurray Date: Sun, 4 Feb 2018 23:48:40 -0800 Subject: [PATCH] Override try_[r]fold for RangeInclusive Because the last item needs special handling, it seems that LLVM has trouble canonicalizing the loops in external iteration. With the override, it becomes obvious that the start==end case exits the loop (as opposed to the one *after* that exiting the loop in external iteration). --- src/libcore/iter/range.rs | 46 ++++++++++++++++++++++++++++++++++++++- src/libcore/tests/iter.rs | 20 +++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/src/libcore/iter/range.rs b/src/libcore/iter/range.rs index 66a76a24df4..3b034efcce1 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}; @@ -397,6 +397,28 @@ 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; + if self.start <= self.end { + loop { + let (x, done) = + if self.start < self.end { + let n = self.start.add_one(); + (mem::replace(&mut self.start, n), false) + } else { + self.end.replace_zero(); + (self.start.replace_one(), true) + }; + accum = f(accum, x)?; + if done { break } + } + } + Try::from_ok(accum) + } } #[unstable(feature = "inclusive_range", reason = "recently added, follows RFC", issue = "28237")] @@ -418,6 +440,28 @@ impl DoubleEndedIterator for ops::RangeInclusive { _ => None, } } + + #[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; + if self.start <= self.end { + loop { + let (x, done) = + if self.start < self.end { + let n = self.end.sub_one(); + (mem::replace(&mut self.end, n), false) + } else { + self.start.replace_one(); + (self.end.replace_zero(), true) + }; + accum = f(accum, x)?; + if done { break } + } + } + Try::from_ok(accum) + } } #[unstable(feature = "fused", issue = "35602")] diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index 8997cf9c6bf..e33a0b6224e 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -1397,6 +1397,26 @@ fn test_range_inclusive_min() { assert_eq!(r.min(), None); } +#[test] +fn test_range_inclusive_folds() { + assert_eq!((1..=10).sum::(), 55); + assert_eq!((1..=10).rev().sum::(), 55); + + let mut it = 40..=50; + assert_eq!(it.try_fold(0, i8::checked_add), None); + assert_eq!(it, 44..=50); + assert_eq!(it.try_rfold(0, i8::checked_add), None); + assert_eq!(it, 44..=47); + + let mut it = 10..=20; + assert_eq!(it.try_fold(0, |a,b| Some(a+b)), Some(165)); + assert_eq!(it, 1..=0); + + let mut it = 10..=20; + assert_eq!(it.try_rfold(0, |a,b| Some(a+b)), Some(165)); + assert_eq!(it, 1..=0); +} + #[test] fn test_repeat() { let mut it = repeat(42);