about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--src/libcollections/hashmap.rs78
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
             }
         }