---

yaml
---
r: 54105
b: refs/heads/dist-snap
c: 89e2578
h: refs/heads/master
i:
  54103: 0cfbdde
v: v3

Author: thestinger

[refs]

@@ -7,7 +7,7 @@ refs/tags/release-0.1: 1f5c5126e96c79d22cb7862f75304136e204f105
 refs/heads/ndm: f3868061cd7988080c30d6d5bf352a5a5fe2460b
 refs/heads/try2: 147ecfdd8221e4a4d4e090486829a06da1e0ca3c
 refs/heads/incoming: 44d4d6de762f3f9aae1fedcf454c66b79b3ad58d
-refs/heads/dist-snap: d77433386be3d5afa201cfd40dbaaa4d20dbbd80
+refs/heads/dist-snap: 89e2578a9d1cd18270b7148a5d5d6b8bee051ac5
 refs/tags/release-0.2: c870d2dffb391e14efb05aa27898f1f6333a9596
 refs/tags/release-0.3: b5f0d0f648d9a6153664837026ba1be43d3e2503
 refs/heads/try3: 9387340aab40a73e8424c48fd42f0c521a4875c0

branches/dist-snap/src/libcore/trie.rs

@@ -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;
@@ -153,6 +153,12 @@ pub impl<T> TrieMap<T> {
     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<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();