std::trie: Fix find_mut for non-present keys

Make TrieMap/TrieSet's find_mut check the key for external nodes.
Without this find_mut sometimes returns a reference to another key when
querying for a non-present key.

1 files changed, 12 insertions, 1 deletions

diff --git a/src/libstd/trie.rs b/src/libstd/trie.rs
index f60d8641aca..98f5baf1e2b 100644
--- a/src/libstd/trie.rs
+++ b/src/libstd/trie.rs
@@ -373,7 +373,8 @@ fn chunk(n: uint, idx: uint) -> uint {
 
 fn find_mut<'r, T>(child: &'r mut Child<T>, key: uint, idx: uint) -> Option<&'r mut T> {
     match *child {
-        External(_, ref mut value) => Some(value),
+        External(stored, ref mut value) if stored == key => Some(value),
+        External(*) => None,
         Internal(ref mut x) => find_mut(&mut x.children[chunk(key, idx)], key, idx + 1),
         Nothing => None
     }
@@ -537,6 +538,16 @@ mod test_map {
     }
 
     #[test]
+    fn test_find_mut_missing() {
+        let mut m = TrieMap::new();
+        assert!(m.find_mut(&0).is_none());
+        assert!(m.insert(1, 12));
+        assert!(m.find_mut(&0).is_none());
+        assert!(m.insert(2, 8));
+        assert!(m.find_mut(&0).is_none());
+    }
+
+    #[test]
     fn test_step() {
         let mut trie = TrieMap::new();
         let n = 300u;