From 2dd6bc6887858f91b846ffa7dbd9f0f49b9f3b87 Mon Sep 17 00:00:00 2001
From: P1start <rewi-github@whanau.org>
Date: Sat, 2 Aug 2014 18:54:38 +1200
Subject: [PATCH] Implement Index for TreeMap

---
 src/libcollections/treemap.rs | 36 +++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/libcollections/treemap.rs b/src/libcollections/treemap.rs
index 6a29a9a75b8e..23f9ae760dcb 100644
--- 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 @@ mod test_treemap {
         }
     }
 
+    #[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)]