Skip to content

Commit 0dbd509

Browse files
committed
auto merge of #9199 : thestinger/rust/range_step, r=cmr
My focus was on getting these to be correct in all cases by handling overflow properly. I'll clean them up and work on the performance later.
2 parents cf7e93f + a18038f commit 0dbd509

File tree

1 file changed

+101
-14
lines changed

1 file changed

+101
-14
lines changed

src/libstd/iter.rs

+101-14
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use cmp;
2121
use num::{Zero, One, Integer, CheckedAdd, CheckedSub, Saturating};
2222
use option::{Option, Some, None};
2323
use ops::{Add, Mul, Sub};
24-
use cmp::Ord;
24+
use cmp::{Eq, Ord};
2525
use clone::Clone;
2626
use uint;
2727
use util;
@@ -1719,7 +1719,21 @@ pub fn count<A>(start: A, step: A) -> Counter<A> {
17191719
Counter{state: start, step: step}
17201720
}
17211721

1722-
/// A range of numbers from [0, N)
1722+
impl<A: Add<A, A> + Clone> Iterator<A> for Counter<A> {
1723+
#[inline]
1724+
fn next(&mut self) -> Option<A> {
1725+
let result = self.state.clone();
1726+
self.state = self.state + self.step;
1727+
Some(result)
1728+
}
1729+
1730+
#[inline]
1731+
fn size_hint(&self) -> (uint, Option<uint>) {
1732+
(uint::max_value, None) // Too bad we can't specify an infinite lower bound
1733+
}
1734+
}
1735+
1736+
/// An iterator over the range [start, stop)
17231737
#[deriving(Clone, DeepClone)]
17241738
pub struct Range<A> {
17251739
priv state: A,
@@ -1749,14 +1763,12 @@ impl<A: Add<A, A> + Ord + Clone> Iterator<A> for Range<A> {
17491763
// Blocked on #8605 Need numeric trait for converting to `Option<uint>`
17501764
}
17511765

1752-
impl<A: Sub<A, A> + Integer + Ord + Clone> DoubleEndedIterator<A> for Range<A> {
1766+
/// `Integer` is required to ensure the range will be the same regardless of
1767+
/// the direction it is consumed.
1768+
impl<A: Integer + Ord + Clone> DoubleEndedIterator<A> for Range<A> {
17531769
#[inline]
17541770
fn next_back(&mut self) -> Option<A> {
17551771
if self.stop > self.state {
1756-
// Integer doesn't technically define this rule, but we're going to assume that every
1757-
// Integer is reachable from every other one by adding or subtracting enough Ones. This
1758-
// seems like a reasonable-enough rule that every Integer should conform to, even if it
1759-
// can't be statically checked.
17601772
self.stop = self.stop - self.one;
17611773
Some(self.stop.clone())
17621774
} else {
@@ -1765,7 +1777,7 @@ impl<A: Sub<A, A> + Integer + Ord + Clone> DoubleEndedIterator<A> for Range<A> {
17651777
}
17661778
}
17671779

1768-
/// A range of numbers from [0, N]
1780+
/// An iterator over the range [start, stop]
17691781
#[deriving(Clone, DeepClone)]
17701782
pub struct RangeInclusive<A> {
17711783
priv range: Range<A>,
@@ -1826,17 +1838,78 @@ impl<A: Sub<A, A> + Integer + Ord + Clone> DoubleEndedIterator<A> for RangeInclu
18261838
}
18271839
}
18281840

1829-
impl<A: Add<A, A> + Clone> Iterator<A> for Counter<A> {
1841+
/// An iterator over the range [start, stop) by `step`. It handles overflow by stopping.
1842+
#[deriving(Clone, DeepClone)]
1843+
pub struct RangeStep<A> {
1844+
priv state: A,
1845+
priv stop: A,
1846+
priv step: A,
1847+
priv rev: bool
1848+
}
1849+
1850+
/// Return an iterator over the range [start, stop) by `step`. It handles overflow by stopping.
1851+
#[inline]
1852+
pub fn range_step<A: CheckedAdd + Ord + Clone + Zero>(start: A, stop: A, step: A) -> RangeStep<A> {
1853+
let rev = step < Zero::zero();
1854+
RangeStep{state: start, stop: stop, step: step, rev: rev}
1855+
}
1856+
1857+
impl<A: CheckedAdd + Ord + Clone> Iterator<A> for RangeStep<A> {
18301858
#[inline]
18311859
fn next(&mut self) -> Option<A> {
1832-
let result = self.state.clone();
1833-
self.state = self.state + self.step;
1834-
Some(result)
1860+
if (self.rev && self.state > self.stop) || self.state < self.stop {
1861+
let result = self.state.clone();
1862+
match self.state.checked_add(&self.step) {
1863+
Some(x) => self.state = x,
1864+
None => self.state = self.stop.clone()
1865+
}
1866+
Some(result)
1867+
} else {
1868+
None
1869+
}
18351870
}
1871+
}
1872+
1873+
/// An iterator over the range [start, stop] by `step`. It handles overflow by stopping.
1874+
#[deriving(Clone, DeepClone)]
1875+
pub struct RangeStepInclusive<A> {
1876+
priv state: A,
1877+
priv stop: A,
1878+
priv step: A,
1879+
priv rev: bool,
1880+
priv done: bool
1881+
}
18361882

1883+
/// Return an iterator over the range [start, stop] by `step`. It handles overflow by stopping.
1884+
#[inline]
1885+
pub fn range_step_inclusive<A: CheckedAdd + Ord + Clone + Zero>(start: A, stop: A,
1886+
step: A) -> RangeStepInclusive<A> {
1887+
let rev = step < Zero::zero();
1888+
RangeStepInclusive{state: start, stop: stop, step: step, rev: rev, done: false}
1889+
}
1890+
1891+
impl<A: CheckedAdd + Ord + Clone + Eq> Iterator<A> for RangeStepInclusive<A> {
18371892
#[inline]
1838-
fn size_hint(&self) -> (uint, Option<uint>) {
1839-
(uint::max_value, None) // Too bad we can't specify an infinite lower bound
1893+
fn next(&mut self) -> Option<A> {
1894+
if !self.done {
1895+
if (self.rev && self.state > self.stop) || self.state < self.stop {
1896+
let result = self.state.clone();
1897+
match self.state.checked_add(&self.step) {
1898+
Some(x) => self.state = x,
1899+
None => self.done = true
1900+
}
1901+
Some(result)
1902+
} else {
1903+
if self.state == self.stop {
1904+
self.done = true;
1905+
Some(self.state.clone())
1906+
} else {
1907+
None
1908+
}
1909+
}
1910+
} else {
1911+
None
1912+
}
18401913
}
18411914
}
18421915

@@ -2649,6 +2722,20 @@ mod tests {
26492722
assert_eq!(range_inclusive(0i, 5).invert().collect::<~[int]>(), ~[5i, 4, 3, 2, 1, 0]);
26502723
}
26512724

2725+
#[test]
2726+
fn test_range_step() {
2727+
assert_eq!(range_step(0i, 20, 5).collect::<~[int]>(), ~[0, 5, 10, 15]);
2728+
assert_eq!(range_step(20i, 0, -5).collect::<~[int]>(), ~[20, 15, 10, 5]);
2729+
assert_eq!(range_step(200u8, 255, 50).collect::<~[u8]>(), ~[200u8, 250]);
2730+
}
2731+
2732+
#[test]
2733+
fn test_range_step_inclusive() {
2734+
assert_eq!(range_step_inclusive(0i, 20, 5).collect::<~[int]>(), ~[0, 5, 10, 15, 20]);
2735+
assert_eq!(range_step_inclusive(20i, 0, -5).collect::<~[int]>(), ~[20, 15, 10, 5, 0]);
2736+
assert_eq!(range_step_inclusive(200u8, 255, 50).collect::<~[u8]>(), ~[200u8, 250]);
2737+
}
2738+
26522739
#[test]
26532740
fn test_reverse() {
26542741
let mut ys = [1, 2, 3, 4, 5];

0 commit comments

Comments
 (0)