iterator: add a `range_inclusive` function

Closes #6242

1 files changed, 52 insertions, 0 deletions

diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index 1d32c5df14e..554913ab5ec 100644
--- a/src/libstd/iterator.rs
+++ b/src/libstd/iterator.rs
@@ -1522,6 +1522,52 @@ impl<A: Sub<A, A> + Integer + Ord + Clone> DoubleEndedIterator<A> for Range<A> {
     }
 }
 
+/// A range of numbers from [0, N]
+#[deriving(Clone, DeepClone)]
+pub struct RangeInclusive<A> {
+    priv range: Range<A>,
+    priv done: bool
+}
+
+/// Return an iterator over the range [start, stop]
+#[inline]
+pub fn range_inclusive<A: Add<A, A> + Ord + Clone + One>(start: A, stop: A) -> RangeInclusive<A> {
+    RangeInclusive{range: range(start, stop), done: false}
+}
+
+impl<A: Add<A, A> + Ord + Clone> Iterator<A> for RangeInclusive<A> {
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        match self.range.next() {
+            Some(x) => Some(x),
+            None => {
+                if self.done {
+                    None
+                } else {
+                    self.done = true;
+                    Some(self.range.stop.clone())
+                }
+            }
+        }
+    }
+}
+
+impl<A: Sub<A, A> + Integer + Ord + Clone> DoubleEndedIterator<A> for RangeInclusive<A> {
+    #[inline]
+    fn next_back(&mut self) -> Option<A> {
+        if self.range.stop > self.range.state {
+            let result = self.range.stop.clone();
+            self.range.stop = self.range.stop - self.range.one;
+            Some(result)
+        } else if self.done {
+            None
+        } else {
+            self.done = true;
+            Some(self.range.stop.clone())
+        }
+    }
+}
+
 impl<A: Add<A, A> + Clone> Iterator<A> for Counter<A> {
     #[inline]
     fn next(&mut self) -> Option<A> {
@@ -2286,4 +2332,10 @@ mod tests {
             fail!("unreachable");
         }
     }
+
+    #[test]
+    fn test_range_inclusive() {
+        assert_eq!(range_inclusive(0i, 5).collect::<~[int]>(), ~[0i, 1, 2, 3, 4, 5]);
+        assert_eq!(range_inclusive(0i, 5).invert().collect::<~[int]>(), ~[5i, 4, 3, 2, 1, 0]);
+    }
 }