diff options
| -rw-r--r-- | src/libcollections/hashmap.rs | 78 |
1 files changed, 57 insertions, 21 deletions
diff --git a/src/libcollections/hashmap.rs b/src/libcollections/hashmap.rs index 3cffea46fec..44312e0eeaf 100644 --- a/src/libcollections/hashmap.rs +++ b/src/libcollections/hashmap.rs @@ -1239,31 +1239,14 @@ impl<K: TotalEq + Hash<S>, V, S, H: Hasher<S>> HashMap<K, V, H> { /// Return the value corresponding to the key in the map, or insert /// and return the value if it doesn't exist. pub fn find_or_insert<'a>(&'a mut self, k: K, v: V) -> &'a mut V { - let hash = self.make_hash(&k); - match self.search_hashed(&hash, &k) { - Some(idx) => { - let (_, v_ref) = self.table.read_mut(&idx); - v_ref - }, - None => self.insert_hashed(hash, k, v) - } + self.mangle(k, v, |_k, a| a, |_k, _v, _a| ()) } /// Return the value corresponding to the key in the map, or create, /// insert, and return a new value if it doesn't exist. pub fn find_or_insert_with<'a>(&'a mut self, k: K, f: |&K| -> V) -> &'a mut V { - let hash = self.make_hash(&k); - match self.search_hashed(&hash, &k) { - Some(idx) => { - let (_, v_ref) = self.table.read_mut(&idx); - v_ref - }, - None => { - let v = f(&k); - self.insert_hashed(hash, k, v) - } - } + self.mangle(k, (), |k, _a| f(k), |_k, _v, _a| ()) } /// Insert a key-value pair into the map if the key is not already present. @@ -1275,12 +1258,65 @@ impl<K: TotalEq + Hash<S>, V, S, H: Hasher<S>> HashMap<K, V, H> { v: V, f: |&K, &mut V|) -> &'a mut V { + self.mangle(k, v, |_k, a| a, |k, v, _a| f(k, v)) + } + + /// Modify and return the value corresponding to the key in the map, or + /// insert and return a new value if it doesn't exist. + /// + /// This method allows for all insertion behaviours of a hashmap; + /// see methods like `insert`, `find_or_insert` and + /// `insert_or_update_with` for less general and more friendly + /// variations of this. + /// + /// # Example + /// + /// ```rust + /// use collections::HashMap; + /// + /// // map some strings to vectors of strings + /// let mut map = HashMap::<StrBuf, Vec<StrBuf>>::new(); + /// map.insert(StrBuf::from_str("a key"), vec![StrBuf::from_str("value")]); + /// map.insert(StrBuf::from_str("z key"), vec![StrBuf::from_str("value")]); + /// + /// let new = vec![StrBuf::from_str("a key"), + /// StrBuf::from_str("b key"), + /// StrBuf::from_str("z key")]; + /// for k in new.move_iter() { + /// map.mangle(k, StrBuf::from_str("new value"), + /// // if the key doesn't exist in the map yet, add it in + /// // the obvious way. + /// |_k, v| vec![v], + /// // if the key does exist either prepend or append this + /// // new value based on the first letter of the key. + /// |key, already, new| { + /// if key.as_slice().starts_with("z") { + /// already.unshift(new); + /// } else { + /// already.push(new); + /// } + /// }); + /// } + /// + /// for (k, v) in map.iter() { + /// println!("{} -> {:?}", *k, *v); + /// } + /// ``` + pub fn mangle<'a, A>(&'a mut self, + k: K, + a: A, + not_found: |&K, A| -> V, + found: |&K, &mut V, A|) + -> &'a mut V { let hash = self.make_hash(&k); match self.search_hashed(&hash, &k) { - None => self.insert_hashed(hash, k, v), + None => { + let v = not_found(&k, a); + self.insert_hashed(hash, k, v) + }, Some(idx) => { let (_, v_ref) = self.table.read_mut(&idx); - f(&k, v_ref); + found(&k, v_ref, a); v_ref } } |
