iterator: Add StepBy::size_hint method

Also fixes `Step::steps_between` implementations by integer types
to correctly handle `by != 1`.

Author: johncf

src/libcore/iter.rs

@@ -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>;
 }
@@ -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)
                 }
@@ -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)
                 }
             }
         }
@@ -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 {

src/libcoretest/iter.rs

@@ -393,6 +393,12 @@ fn test_iterator_size_hint() {
     assert_eq!(vi.clone().filter(|_| false).size_hint(), (0, Some(10)));
     assert_eq!(vi.clone().map(|&i| i+1).size_hint(), (10, Some(10)));
     assert_eq!(vi.filter_map(|_| Some(0)).size_hint(), (0, Some(10)));
+
+    assert_eq!((0..20).step_by(1).size_hint(), (20, Some(20)));
+    assert_eq!((0..20).step_by(5).size_hint(), (4, Some(4)));
+    assert_eq!((0..20).step_by(6).size_hint(), (4, Some(4)));
+    assert_eq!((20..0).step_by(-6).size_hint(), (4, Some(4)));
+    assert_eq!((0..20).step_by(21).size_hint(), (1, Some(1)));
 }
 
 #[test]
@@ -788,6 +794,20 @@ fn test_range_step() {
     assert_eq!((200..200).step_by(1).collect::<Vec<isize>>(), []);
 }
 
+#[test]
+fn test_steps_between() {
+    assert_eq!(Step::steps_between(&0, &20, &1), Some(20));
+    assert_eq!(Step::steps_between(&0, &20, &21), Some(1));
+    assert_eq!(Step::steps_between(&0, &20, &5), Some(4));
+    assert_eq!(Step::steps_between(&20, &0, &-5), Some(4));
+    assert_eq!(Step::steps_between(&20, &-5, &1), Some(0));
+    assert_eq!(Step::steps_between(&20, &20, &1), Some(0));
+    assert_eq!(Step::steps_between(&0, &1, &0), None);
+    assert_eq!(Step::steps_between(&std::i8::MAX, &std::i8::MIN, &std::i8::MIN), Some(2));
+    assert_eq!(Step::steps_between(&std::i16::MIN, &std::i16::MAX, &std::i16::MAX), Some(3));
+    assert_eq!(Step::steps_between(&std::isize::MIN, &std::isize::MAX, &1), Some(std::usize::MAX));
+}
+
 #[test]
 fn test_reverse() {
     let mut ys = [1, 2, 3, 4, 5];