Implement Index for TreeMap

diff --git a/src/libcollections/treemap.rs b/src/libcollections/treemap.rs
index 6a29a9a75b8e10bf9d70ac3e63c7df65d3feed33..23f9ae760dcb2dbf4ef9978d10d67baaa771a3e4 100644 (file)
--- a/src/libcollections/treemap.rs
+++ b/src/libcollections/treemap.rs
@@ -237,6 +237,20 @@ impl<K: Ord, V> Default for TreeMap<K,V> {
     fn default() -> TreeMap<K, V> { TreeMap::new() }
 }
 
+impl<K: Ord, V> Index<K, V> for TreeMap<K, V> {
+    #[inline]
+    fn index<'a>(&'a self, i: &K) -> &'a V {
+        self.find(i).expect("no entry found for key")
+    }
+}
+
+/*impl<K: Ord, V> IndexMut<K, V> for TreeMap<K, V> {
+    #[inline]
+    fn index_mut<'a>(&'a mut self, i: &K) -> &'a mut V {
+        self.find_mut(i).expect("no entry found for key")
+    }
+}*/
+
 impl<K: Ord, V> TreeMap<K, V> {
     /// Create an empty `TreeMap`.
     ///
@@ -2131,6 +2145,28 @@ fn test_from_iter() {
         }
     }
 
+    #[test]
+    fn test_index() {
+        let mut map: TreeMap<int, int> = TreeMap::new();
+
+        map.insert(1, 2);
+        map.insert(2, 1);
+        map.insert(3, 4);
+
+        assert_eq!(map[2], 1);
+    }
+
+    #[test]
+    #[should_fail]
+    fn test_index_nonexistent() {
+        let mut map: TreeMap<int, int> = TreeMap::new();
+
+        map.insert(1, 2);
+        map.insert(2, 1);
+        map.insert(3, 4);
+
+        map[4];
+    }
 }
 
 #[cfg(test)]