std: add an external iterator for mutating vec elements

1 files changed, 44 insertions, 0 deletions

diff --git a/src/libstd/vec.rs b/src/libstd/vec.rs
index f078cd3bda3..58b0193d300 100644
--- a/src/libstd/vec.rs
+++ b/src/libstd/vec.rs
@@ -2427,6 +2427,7 @@ impl<T:Eq> OwnedEqVector<T> for ~[T] {
 #[allow(missing_doc)]
 pub trait MutableVector<'self, T> {
     fn mut_slice(self, start: uint, end: uint) -> &'self mut [T];
+    fn mut_iter(self) -> MutVecIterator<'self, T>;
 
     unsafe fn unsafe_mut_ref(&self, index: uint) -> *mut T;
     unsafe fn unsafe_set(&self, index: uint, val: T);
@@ -2438,6 +2439,15 @@ impl<'self,T> MutableVector<'self, T> for &'self mut [T] {
         mut_slice(self, start, end)
     }
 
+    #[inline]
+    fn mut_iter(self) -> MutVecIterator<'self, T> {
+        unsafe {
+            let p = vec::raw::to_mut_ptr(self);
+            MutVecIterator{ptr: p, end: p.offset(self.len()),
+                           lifetime: cast::transmute(p)}
+        }
+    }
+
     #[inline(always)]
     unsafe fn unsafe_mut_ref(&self, index: uint) -> *mut T {
         let pair_ptr: &(*mut T, uint) = transmute(self);
@@ -2962,6 +2972,30 @@ impl<'self, T> Iterator<&'self T> for VecIterator<'self, T> {
     }
 }
 
+/// An external iterator for vectors with the possibility of mutating
+/// elements. (use with the std::iterator module)
+pub struct MutVecIterator<'self, T> {
+    priv ptr: *mut T,
+    priv end: *mut T,
+    priv lifetime: &'self mut T // FIXME: #5922
+}
+
+// could be implemented with &[T] with .slice(), but this avoids bounds checks
+impl<'self, T> Iterator<&'self mut T> for MutVecIterator<'self, T> {
+    #[inline]
+    fn next(&mut self) -> Option<&'self mut T> {
+        unsafe {
+            if self.ptr == self.end {
+                None
+            } else {
+                let old = self.ptr;
+                self.ptr = self.ptr.offset(1);
+                Some(cast::transmute(old))
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use option::{None, Option, Some};
@@ -4670,6 +4704,16 @@ mod tests {
     }
 
     #[test]
+    fn test_mut_iterator() {
+        use iterator::*;
+        let mut xs = [1, 2, 3, 4, 5];
+        for xs.mut_iter().advance |x| {
+            *x += 1;
+        }
+        assert_eq!(xs, [2, 3, 4, 5, 6])
+    }
+
+    #[test]
     fn test_reverse_part() {
         let mut values = [1,2,3,4,5];
         reverse_part(values,1,4);