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: 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 Map for TreeMap { 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> = &self.root; loop { @@ -189,6 +189,12 @@ pub impl TreeMap { 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(node: &mut ~TreeNode) { } } -fn insert(node: &mut Option<~TreeNode>, key: K, - value: V) -> bool { +fn find_mut(node: &'r mut Option<~TreeNode>, 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(node: &mut Option<~TreeNode>, 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();