Skip to content

Commit

Permalink
Override try_[r]fold for RangeInclusive
Browse files Browse the repository at this point in the history
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).
  • Loading branch information
scottmcm committed Feb 5, 2018
1 parent b8f2674 commit 1b1e887
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 1 deletion.
46 changes: 45 additions & 1 deletion src/libcore/iter/range.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -397,6 +397,28 @@ impl<A: Step> Iterator for ops::RangeInclusive<A> {
fn max(mut self) -> Option<A> {
self.next_back()
}

#[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.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")]
Expand All @@ -418,6 +440,28 @@ impl<A: Step> DoubleEndedIterator for ops::RangeInclusive<A> {
_ => None,
}
}

#[inline]
fn try_rfold<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.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")]
Expand Down
20 changes: 20 additions & 0 deletions src/libcore/tests/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<i32>(), 55);
assert_eq!((1..=10).rev().sum::<i32>(), 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);
Expand Down

0 comments on commit 1b1e887

Please sign in to comment.