Skip to content

Commit

Permalink
iterator: Add StepBy::size_hint method
Browse files Browse the repository at this point in the history
Fixes `Step::steps_between` implementations by integer types
to correctly handle `by != 1`.
  • Loading branch information
johncf committed May 1, 2015
1 parent c48b499 commit 2a8fc9b
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 16 deletions.
56 changes: 41 additions & 15 deletions src/libcore/iter.rs
Expand Up @@ -2402,14 +2402,10 @@ pub trait Step: PartialOrd {
/// Steps `self` if possible.
fn step(&self, by: &Self) -> Option<Self>;

/// Returns the number of steps between two step objects.
/// Returns the number of steps between two step objects. The count is
/// inclusive of `start` and exclusive of `end`.
///
/// `start` should always be less than `end`, so the result should never
/// be negative.
///
/// `by` must be > 0.
///
/// Returns `None` if it is not possible to calculate steps_between
/// Returns `None` if it is not possible to calculate `steps_between`
/// without overflow.
fn steps_between(start: &Self, end: &Self, by: &Self) -> Option<usize>;
}
Expand All @@ -2424,9 +2420,16 @@ macro_rules! step_impl_unsigned {
#[inline]
#[allow(trivial_numeric_casts)]
fn steps_between(start: &$t, end: &$t, by: &$t) -> Option<usize> {
if *start <= *end {
if *by == 0 { return None; }
if *start < *end {
// Note: We assume $t <= usize here
Some((*end - *start) as usize / (*by as usize))
let diff = (*end - *start) as usize;
let by = *by as usize;
if diff % by > 0 {
Some(diff / by + 1)
} else {
Some(diff / by)
}
} else {
Some(0)
}
Expand All @@ -2444,16 +2447,29 @@ macro_rules! step_impl_signed {
#[inline]
#[allow(trivial_numeric_casts)]
fn steps_between(start: &$t, end: &$t, by: &$t) -> Option<usize> {
if *start <= *end {
if *by == 0 { return None; }
let mut diff: usize;
let mut by_u: usize;
if *by > 0 {
if *start >= *end {
return Some(0);
}
// Note: We assume $t <= isize here
// Use .wrapping_sub and cast to usize to compute the
// difference that may not fit inside the range of isize.
Some(
((*end as isize).wrapping_sub(*start as isize) as usize
/ (*by as usize))
)
diff = (*end as isize).wrapping_sub(*start as isize) as usize;
by_u = *by as usize;
} else {
Some(0)
if *start <= *end {
return Some(0);
}
diff = (*start as isize).wrapping_sub(*end as isize) as usize;
by_u = (*by as isize).wrapping_mul(-1) as usize;
}
if diff % by_u > 0 {
Some(diff / by_u + 1)
} else {
Some(diff / by_u)
}
}
}
Expand Down Expand Up @@ -2675,6 +2691,16 @@ impl<A: Step + Zero + Clone> Iterator for StepBy<A, ops::Range<A>> {
None
}
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
match Step::steps_between(&self.range.start,
&self.range.end,
&self.step_by) {
Some(hint) => (hint, Some(hint)),
None => (0, None)
}
}
}

macro_rules! range_exact_iter_impl {
Expand Down
14 changes: 13 additions & 1 deletion src/libcoretest/iter.rs
Expand Up @@ -11,7 +11,7 @@
use core::iter::*;
use core::iter::order::*;
use core::iter::MinMaxResult::*;
use core::isize;
use core::{i8, i16, isize};
use core::usize;
use core::cmp;

Expand Down Expand Up @@ -786,6 +786,18 @@ fn test_range_step() {
assert_eq!((200..255).step_by(50).collect::<Vec<u8>>(), [200, 250]);
assert_eq!((200..-5).step_by(1).collect::<Vec<isize>>(), []);
assert_eq!((200..200).step_by(1).collect::<Vec<isize>>(), []);

assert_eq!((0..20).step_by(1).size_hint(), (20, Some(20)));
assert_eq!((0..20).step_by(21).size_hint(), (1, Some(1)));
assert_eq!((0..20).step_by(5).size_hint(), (4, Some(4)));
assert_eq!((20..0).step_by(-5).size_hint(), (4, Some(4)));
assert_eq!((20..0).step_by(-6).size_hint(), (4, Some(4)));
assert_eq!((20..-5).step_by(1).size_hint(), (0, Some(0)));
assert_eq!((20..20).step_by(1).size_hint(), (0, Some(0)));
assert_eq!((0..1).step_by(0).size_hint(), (0, None));
assert_eq!((i8::MAX..i8::MIN).step_by(i8::MIN).size_hint(), (2, Some(2)));
assert_eq!((i16::MIN..i16::MAX).step_by(i16::MAX).size_hint(), (3, Some(3)));
assert_eq!((isize::MIN..isize::MAX).step_by(1).size_hint(), (usize::MAX, Some(usize::MAX)));
}

#[test]
Expand Down

0 comments on commit 2a8fc9b

Please sign in to comment.