add find_mut method to the Map trait

src/libcore/container.rs

@@ -38,9 +38,12 @@ pub trait Map<K, V>: Mutable {
     /// Iterate over the map and mutate the contained values
     fn mutate_values(&mut self, f: &fn(&K, &mut V) -> bool);
 
-    /// Return the value corresponding to the key in the map
+    /// Return a reference to the value corresponding to the key
     fn find(&self, key: &K) -> Option<&'self V>;
 
+    /// Return a mutable reference to the value corresponding to the key
+    fn find_mut(&mut self, key: &K) -> Option<&'self mut V>;
+
     /// Insert a key-value pair into the map. An existing value for a
     /// key is replaced by the new value. Return true if the key did
     /// not already exist in the map.

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 @@ pub mod linear {
             }
         }
 
+        #[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,14 +347,25 @@ pub mod linear {
             }
         }
 
-        /// 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)),
                 TableFull | FoundHole(_) => None,
             }
         }
 
+        /// 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)))
+            }
+        }
+
         /// Insert a key-value pair into the map. An existing value for a
         /// key is replaced by the new value. Return true if the key did
         /// not already exist in the map.
@@ -655,6 +675,19 @@ pub mod linear {
             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();

src/libcore/trie.rs

@@ -8,7 +8,7 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-//! A radix trie for storing integers in sorted order
+//! An ordered map and set for integer keys implemented as a radix trie
 
 use prelude::*;
 
@@ -90,7 +90,7 @@ impl<T> Map<uint, T> for TrieMap<T> {
         self.root.mutate_values(f);
     }
 
-    /// Return the value corresponding to the key in the map
+    /// Return a reference to the value corresponding to the key
     #[inline(hint)]
     fn find(&self, key: &uint) -> Option<&'self T> {
         let mut node: &'self TrieNode<T> = &self.root;
@@ -111,6 +111,12 @@ impl<T> Map<uint, T> for TrieMap<T> {
         }
     }
 
+    /// Return a mutable reference to the value corresponding to the key
+    #[inline(always)]
+    fn find_mut(&mut self, key: &uint) -> Option<&'self mut T> {
+        find_mut(&mut self.root.children[chunk(*key, 0)], *key, 1)
+    }
+
     /// Insert a key-value pair into the map. An existing value for a
     /// key is replaced by the new value. Return true if the key did
     /// not already exist in the map.
@@ -276,6 +282,17 @@ fn chunk(n: uint, idx: uint) -> uint {
     (n >> sh) & MASK
 }
 
+fn find_mut<T>(child: &'r mut Child<T>, key: uint, idx: uint) -> Option<&'r mut T> {
+    unsafe { // FIXME(#4903)---requires flow-sensitive borrow checker
+        (match *child {
+            External(_, ref value) => Some(cast::transmute_mut(value)),
+            Internal(ref x) => find_mut(cast::transmute_mut(&x.children[chunk(key, idx)]),
+                                        key, idx + 1),
+            Nothing => None
+        }).map_consume(|x| cast::transmute_mut_region(x))
+    }
+}
+
 fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
              idx: uint) -> bool {
     let mut tmp = Nothing;
@@ -357,8 +374,22 @@ pub fn check_integrity<T>(trie: &TrieNode<T>) {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use core::option::{Some, None};
     use uint;
 
+    #[test]
+    fn test_find_mut() {
+        let mut m = TrieMap::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]
     fn test_step() {
         let mut trie = TrieMap::new();

src/libstd/smallintmap.rs

@@ -86,7 +86,7 @@ impl<V> Map<uint, V> for SmallIntMap<V> {
         self.each(|&(_, v)| blk(v))
     }
 
-    /// Visit all key-value pairs in order
+    /// Iterate over the map and mutate the contained values
     fn mutate_values(&mut self, it: &fn(&uint, &'self mut V) -> bool) {
         for uint::range(0, self.v.len()) |i| {
             match self.v[i] {
@@ -96,7 +96,7 @@ impl<V> Map<uint, V> for SmallIntMap<V> {
         }
     }
 
-    /// Iterate over the map and mutate the contained values
+    /// Return a reference to the value corresponding to the key
     fn find(&self, key: &uint) -> Option<&'self V> {
         if *key < self.v.len() {
             match self.v[*key] {
@@ -108,6 +108,18 @@ impl<V> Map<uint, V> for SmallIntMap<V> {
         }
     }
 
+    /// Return a mutable reference to the value corresponding to the key
+    fn find_mut(&mut self, key: &uint) -> Option<&'self mut V> {
+        if *key < self.v.len() {
+            match self.v[*key] {
+              Some(ref mut value) => Some(value),
+              None => None
+            }
+        } else {
+            None
+        }
+    }
+
     /// Insert a key-value pair into the map. An existing value for a
     /// key is replaced by the new value. Return true if the key did
     /// not already exist in the map.
@@ -160,6 +172,20 @@ pub impl<V:Copy> SmallIntMap<V> {
 #[cfg(test)]
 mod tests {
     use super::SmallIntMap;
+    use core::prelude::*;
+
+    #[test]
+    fn test_find_mut() {
+        let mut m = SmallIntMap::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]
     fn test_len() {

src/libstd/treemap.rs

@@ -135,7 +135,7 @@ impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
         mutate_values(&mut self.root, f);
     }
 
-    /// Return the value corresponding to the key in the map
+    /// Return a reference to the value corresponding to the key
     fn find(&self, key: &K) -> Option<&'self V> {
         let mut current: &'self Option<~TreeNode<K, V>> = &self.root;
         loop {
@@ -152,6 +152,12 @@ impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
         }
     }
 
+    /// Return a mutable reference to the value corresponding to the key
+    #[inline(always)]
+    fn find_mut(&mut self, key: &K) -> Option<&'self mut V> {
+        find_mut(&mut self.root, key)
+    }
+
     /// Insert a key-value pair into the map. An existing value for a
     /// key is replaced by the new value. Return true if the key did
     /// not already exist in the map.
@@ -584,8 +590,20 @@ fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
     }
 }
 
-fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K,
-                          value: V) -> bool {
+fn find_mut<K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>, key: &K) -> Option<&'r mut V> {
+    match *node {
+      Some(ref mut x) => {
+        match key.cmp(&x.key) {
+          Less => find_mut(&mut x.left, key),
+          Greater => find_mut(&mut x.right, key),
+          Equal => Some(&mut x.value),
+        }
+      }
+      None => None
+    }
+}
+
+fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K, value: V) -> bool {
     match *node {
       Some(ref mut save) => {
         match key.cmp(&save.key) {
@@ -716,6 +734,19 @@ mod test_treemap {
         fail_unless!(m.find(&2) == None);
     }
 
+    #[test]
+    fn test_find_mut() {
+        let mut m = TreeMap::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]
     fn insert_replace() {
         let mut m = TreeMap::new();

src/test/run-pass/class-impl-very-parameterized-trait.rs

@@ -98,6 +98,8 @@ impl<T> Map<int, T> for cat<T> {
         }
     }
 
+    fn find_mut(&mut self, k: &int) -> Option<&'self mut T> { fail!() }
+
     fn remove(&mut self, k: &int) -> bool {
         if self.find(k).is_some() {
             self.meows -= *k; true