diff --git a/src/libstd/iter.rs b/src/libstd/iter.rs index faa2d715df4..22e8a1defbd 100644 --- a/src/libstd/iter.rs +++ b/src/libstd/iter.rs @@ -21,7 +21,7 @@ use cmp; use num::{Zero, One, Integer, CheckedAdd, CheckedSub, Saturating}; use option::{Option, Some, None}; use ops::{Add, Mul, Sub}; -use cmp::Ord; +use cmp::{Eq, Ord}; use clone::Clone; use uint; use util; @@ -1838,6 +1838,81 @@ impl + Integer + Ord + Clone> DoubleEndedIterator for RangeInclu } } +/// An iterator over the range [start, stop) by `step`. It handles overflow by stopping. +#[deriving(Clone, DeepClone)] +pub struct RangeStep { + priv state: A, + priv stop: A, + priv step: A, + priv rev: bool +} + +/// Return an iterator over the range [start, stop) by `step`. It handles overflow by stopping. +#[inline] +pub fn range_step(start: A, stop: A, step: A) -> RangeStep { + let rev = step < Zero::zero(); + RangeStep{state: start, stop: stop, step: step, rev: rev} +} + +impl Iterator for RangeStep { + #[inline] + fn next(&mut self) -> Option { + if (self.rev && self.state > self.stop) || self.state < self.stop { + let result = self.state.clone(); + match self.state.checked_add(&self.step) { + Some(x) => self.state = x, + None => self.state = self.stop.clone() + } + Some(result) + } else { + None + } + } +} + +/// An iterator over the range [start, stop] by `step`. It handles overflow by stopping. +#[deriving(Clone, DeepClone)] +pub struct RangeStepInclusive { + priv state: A, + priv stop: A, + priv step: A, + priv rev: bool, + priv done: bool +} + +/// Return an iterator over the range [start, stop] by `step`. It handles overflow by stopping. +#[inline] +pub fn range_step_inclusive(start: A, stop: A, + step: A) -> RangeStepInclusive { + let rev = step < Zero::zero(); + RangeStepInclusive{state: start, stop: stop, step: step, rev: rev, done: false} +} + +impl Iterator for RangeStepInclusive { + #[inline] + fn next(&mut self) -> Option { + if !self.done { + if (self.rev && self.state > self.stop) || self.state < self.stop { + let result = self.state.clone(); + match self.state.checked_add(&self.step) { + Some(x) => self.state = x, + None => self.done = true + } + Some(result) + } else { + if self.state == self.stop { + self.done = true; + Some(self.state.clone()) + } else { + None + } + } + } else { + None + } + } +} + /// An iterator that repeats an element endlessly #[deriving(Clone, DeepClone)] pub struct Repeat { @@ -2647,6 +2722,20 @@ mod tests { assert_eq!(range_inclusive(0i, 5).invert().collect::<~[int]>(), ~[5i, 4, 3, 2, 1, 0]); } + #[test] + fn test_range_step() { + assert_eq!(range_step(0i, 20, 5).collect::<~[int]>(), ~[0, 5, 10, 15]); + assert_eq!(range_step(20i, 0, -5).collect::<~[int]>(), ~[20, 15, 10, 5]); + assert_eq!(range_step(200u8, 255, 50).collect::<~[u8]>(), ~[200u8, 250]); + } + + #[test] + fn test_range_step_inclusive() { + assert_eq!(range_step_inclusive(0i, 20, 5).collect::<~[int]>(), ~[0, 5, 10, 15, 20]); + assert_eq!(range_step_inclusive(20i, 0, -5).collect::<~[int]>(), ~[20, 15, 10, 5, 0]); + assert_eq!(range_step_inclusive(200u8, 255, 50).collect::<~[u8]>(), ~[200u8, 250]); + } + #[test] fn test_reverse() { let mut ys = [1, 2, 3, 4, 5];