treemap: add a find_mut method

2 files changed, 38 insertions, 4 deletions

diff --git a/src/libcore/container.rs b/src/libcore/container.rs
index 1cda04ee06e..9eba2c7105c 100644
--- a/src/libcore/container.rs
+++ b/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.
diff --git a/src/libstd/treemap.rs b/src/libstd/treemap.rs
index 1da1edae7d2..b77037ba3ad 100644
--- a/src/libstd/treemap.rs
+++ b/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 {
@@ -189,6 +189,12 @@ pub impl<K: TotalOrd, V> TreeMap<K, V> {
     fn iter(&self) -> TreeMapIterator<'self, K, V> {
         TreeMapIterator{stack: ~[], node: &self.root}
     }
+
+    /// 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)
+    }
 }
 
 /// Lazy forward iterator over a 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) {
@@ -717,6 +735,19 @@ mod test_treemap {
     }
 
     #[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();
         fail_unless!(m.insert(5, 2));