From d77433386be3d5afa201cfd40dbaaa4d20dbbd80 Mon Sep 17 00:00:00 2001
From: Daniel Micay <danielmicay@gmail.com>
Date: Sun, 24 Mar 2013 19:07:36 -0400
Subject: [PATCH] hashmap: add find_mut method

---
 src/libcore/hashmap.rs | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/libcore/hashmap.rs b/src/libcore/hashmap.rs
index f5a97bdaca3..b980082a88c 100644
--- a/src/libcore/hashmap.rs
+++ b/src/libcore/hashmap.rs
@@ -24,6 +24,7 @@ pub mod linear {
     use rand;
     use uint;
     use vec;
+    use util::unreachable;
 
     static INITIAL_CAPACITY: uint = 32u; // 2^5
 
@@ -192,6 +193,14 @@ fn value_for_bucket(&self, idx: uint) -> &'self V {
             }
         }
 
+        #[inline(always)]
+        fn mut_value_for_bucket(&mut self, idx: uint) -> &'self mut V {
+            match self.buckets[idx] {
+                Some(ref mut bkt) => &mut bkt.value,
+                None => unreachable()
+            }
+        }
+
         /// Inserts the key value pair into the buckets.
         /// Assumes that there will be a bucket.
         /// True if there was no previous entry with that key
@@ -338,7 +347,7 @@ fn mutate_values(&mut self, blk: &fn(&'self K,
             }
         }
 
-        /// Return the value corresponding to the key in the map
+        /// Return a reference to the value corresponding to the key
         fn find(&self, k: &K) -> Option<&'self V> {
             match self.bucket_for_key(k) {
                 FoundEntry(idx) => Some(self.value_for_bucket(idx)),
@@ -410,6 +419,17 @@ fn swap(&mut self, k: K, v: V) -> Option<V> {
             old_value
         }
 
+        /// Return a mutable reference to the value corresponding to the key
+        fn find_mut(&mut self, k: &K) -> Option<&'self mut V> {
+            let idx = match self.bucket_for_key(k) {
+                FoundEntry(idx) => idx,
+                TableFull | FoundHole(_) => return None
+            };
+            unsafe {  // FIXME(#4903)---requires flow-sensitive borrow checker
+                Some(::cast::transmute_mut_region(self.mut_value_for_bucket(idx)))
+            }
+        }
+
         /// Return the value corresponding to the key in the map, or insert
         /// and return the value if it doesn't exist.
         fn find_or_insert(&mut self, k: K, v: V) -> &'self V {
@@ -655,6 +675,19 @@ pub fn test_insert() {
             fail_unless!(*m.get(&2) == 4);
         }
 
+        #[test]
+        fn test_find_mut() {
+            let mut m = LinearMap::new();
+            fail_unless!(m.insert(1, 12));
+            fail_unless!(m.insert(2, 8));
+            fail_unless!(m.insert(5, 14));
+            let new = 100;
+            match m.find_mut(&5) {
+                None => fail!(), Some(x) => *x = new
+            }
+            assert_eq!(m.find(&5), Some(&new));
+        }
+
         #[test]
         pub fn test_insert_overwrite() {
             let mut m = LinearMap::new();