diff --git a/src/libcore/trie.rs b/src/libcore/trie.rs index 40ef5fee47a..007bafcd03d 100644 --- a/src/libcore/trie.rs +++ b/src/libcore/trie.rs @@ -90,7 +90,7 @@ impl Map for TrieMap { 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 = &self.root; @@ -153,6 +153,12 @@ pub impl TrieMap { fn each_value_reverse(&self, f: &fn(&T) -> bool) { self.each_reverse(|&(_, v)| f(v)) } + + /// 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) + } } pub struct TrieSet { @@ -276,6 +282,17 @@ fn chunk(n: uint, idx: uint) -> uint { (n >> sh) & MASK } +fn find_mut(child: &'r mut Child, 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(count: &mut uint, child: &mut Child, key: uint, value: T, idx: uint) -> bool { let mut tmp = Nothing; @@ -357,8 +374,22 @@ pub fn check_integrity(trie: &TrieNode) { #[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();