add a RandomAccessIterator trait

2 files changed, 117 insertions, 2 deletions

diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index 198e63f83c6..2c971f0c9ce 100644
--- a/src/libstd/iterator.rs
+++ b/src/libstd/iterator.rs
@@ -53,6 +53,16 @@ pub trait DoubleEndedIterator<A>: Iterator<A> {
     fn next_back(&mut self) -> Option<A>;
 }
 
+/// An object implementing random access indexing by `uint`
+pub trait RandomAccessIterator<A> {
+    /// Return the number of indexable elements. At most `std::uint::max_value`
+    /// elements are indexable, even if the iterator represents a longer range.
+    fn indexable(&self) -> uint;
+
+    /// Return an element at an index
+    fn idx(&self, index: uint) -> Option<A>;
+}
+
 /// Iterator adaptors provided for every `DoubleEndedIterator` implementation.
 ///
 /// In the future these will be default methods instead of a utility trait.
@@ -836,6 +846,30 @@ for ChainIterator<A, T, U> {
     }
 }
 
+impl<A, T: RandomAccessIterator<A>, U: RandomAccessIterator<A>> RandomAccessIterator<A>
+for ChainIterator<A, T, U> {
+    #[inline]
+    fn indexable(&self) -> uint {
+        let (a, b) = (self.a.indexable(), self.b.indexable());
+        let total = a + b;
+        if total < a || total < b {
+            uint::max_value
+        } else {
+            total
+        }
+    }
+
+    #[inline]
+    fn idx(&self, index: uint) -> Option<A> {
+        let len = self.a.indexable();
+        if index < len {
+            self.a.idx(index)
+        } else {
+            self.b.idx(index - len)
+        }
+    }
+}
+
 /// An iterator which iterates two other iterators simultaneously
 // FIXME #6967: Dummy A & B parameters to get around type inference bug
 #[deriving(Clone)]
@@ -1718,4 +1752,23 @@ mod tests {
         assert_eq!(it.next_back().unwrap(), &7)
         assert_eq!(it.next_back(), None)
     }
+
+    #[test]
+    fn test_random_access_chain() {
+        let xs = [1, 2, 3, 4, 5];
+        let ys = ~[7, 9, 11];
+        let mut it = xs.iter().chain_(ys.iter());
+        assert_eq!(it.idx(0).unwrap(), &1);
+        assert_eq!(it.idx(5).unwrap(), &7);
+        assert_eq!(it.idx(7).unwrap(), &11);
+        assert!(it.idx(8).is_none());
+
+        it.next();
+        it.next();
+        it.next_back();
+
+        assert_eq!(it.idx(0).unwrap(), &3);
+        assert_eq!(it.idx(4).unwrap(), &9);
+        assert!(it.idx(6).is_none());
+    }
 }
diff --git a/src/libstd/vec.rs b/src/libstd/vec.rs
index baeb87e51b9..9c66ee5eae4 100644
--- a/src/libstd/vec.rs
+++ b/src/libstd/vec.rs
@@ -2116,8 +2116,7 @@ macro_rules! iterator {
 
             #[inline]
             fn size_hint(&self) -> (uint, Option<uint>) {
-                let diff = (self.end as uint) - (self.ptr as uint);
-                let exact = diff / sys::nonzero_size_of::<$elem>();
+                let exact = self.indexable();
                 (exact, Some(exact))
             }
         }
@@ -2145,6 +2144,28 @@ macro_rules! double_ended_iterator {
     }
 }
 
+macro_rules! random_access_iterator {
+    (impl $name:ident -> $elem:ty) => {
+        impl<'self, T> RandomAccessIterator<$elem> for $name<'self, T> {
+            #[inline]
+            fn indexable(&self) -> uint {
+                let diff = (self.end as uint) - (self.ptr as uint);
+                diff / sys::nonzero_size_of::<T>()
+            }
+
+            fn idx(&self, index: uint) -> Option<$elem> {
+                unsafe {
+                    if index < self.indexable() {
+                        cast::transmute(self.ptr.offset(index))
+                    } else {
+                        None
+                    }
+                }
+            }
+        }
+    }
+}
+
 //iterator!{struct VecIterator -> *T, &'self T}
 /// An iterator for iterating over a vector.
 pub struct VecIterator<'self, T> {
@@ -2154,6 +2175,7 @@ pub struct VecIterator<'self, T> {
 }
 iterator!{impl VecIterator -> &'self T}
 double_ended_iterator!{impl VecIterator -> &'self T}
+random_access_iterator!{impl VecIterator -> &'self T}
 pub type VecRevIterator<'self, T> = InvertIterator<&'self T, VecIterator<'self, T>>;
 
 impl<'self, T> Clone for VecIterator<'self, T> {
@@ -2169,6 +2191,7 @@ pub struct VecMutIterator<'self, T> {
 }
 iterator!{impl VecMutIterator -> &'self mut T}
 double_ended_iterator!{impl VecMutIterator -> &'self mut T}
+random_access_iterator!{impl VecMutIterator -> &'self mut T}
 pub type VecMutRevIterator<'self, T> = InvertIterator<&'self mut T, VecMutIterator<'self, T>>;
 
 /// An iterator that moves out of a vector.
@@ -3109,6 +3132,45 @@ mod tests {
     }
 
     #[test]
+    fn test_random_access_iterator() {
+        use iterator::*;
+        let xs = [1, 2, 5, 10, 11];
+        let mut it = xs.iter();
+
+        assert_eq!(it.indexable(), 5);
+        assert_eq!(it.idx(0).unwrap(), &1);
+        assert_eq!(it.idx(2).unwrap(), &5);
+        assert_eq!(it.idx(4).unwrap(), &11);
+        assert!(it.idx(5).is_none());
+
+        assert_eq!(it.next().unwrap(), &1);
+        assert_eq!(it.indexable(), 4);
+        assert_eq!(it.idx(0).unwrap(), &2);
+        assert_eq!(it.idx(3).unwrap(), &11);
+        assert!(it.idx(4).is_none());
+
+        assert_eq!(it.next().unwrap(), &2);
+        assert_eq!(it.indexable(), 3);
+        assert_eq!(it.idx(1).unwrap(), &10);
+        assert!(it.idx(3).is_none());
+
+        assert_eq!(it.next().unwrap(), &5);
+        assert_eq!(it.indexable(), 2);
+        assert_eq!(it.idx(1).unwrap(), &11);
+
+        assert_eq!(it.next().unwrap(), &10);
+        assert_eq!(it.indexable(), 1);
+        assert_eq!(it.idx(0).unwrap(), &11);
+        assert!(it.idx(1).is_none());
+
+        assert_eq!(it.next().unwrap(), &11);
+        assert_eq!(it.indexable(), 0);
+        assert!(it.idx(0).is_none());
+
+        assert!(it.next().is_none());
+    }
+
+    #[test]
     fn test_iter_size_hints() {
         use iterator::*;
         let mut xs = [1, 2, 5, 10, 11];