Implement HashMap.find_with_or_insert_with()

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.find_with_or_insert_with(k, v, |_k, _v, _a| (), |_k, a| 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.find_with_or_insert_with(k, (), |_k, _v, _a| (), |k, _a| f(k))
     }
 
     /// Insert a key-value pair into the map if the key is not already present.
@@ -1275,12 +1258,66 @@ impl<K: TotalEq + Hash<S>, V, S, H: Hasher<S>> HashMap<K, V, H> {
                                  v: V,
                                  f: |&K, &mut V|)
                                  -> &'a mut V {
+        self.find_with_or_insert_with(k, v, |k, v, _a| f(k, v), |_k, a| a)
+    }
+
+    /// 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::new();
+    /// map.insert("a key", vec!["value"]);
+    /// map.insert("z key", vec!["value"]);
+    ///
+    /// let new = vec!["a key", "b key", "z key"];
+    ///
+    /// for k in new.move_iter() {
+    ///     map.find_with_or_insert_with(
+    ///         k, "new value",
+    ///         // 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);
+    ///             }
+    ///         },
+    ///         // if the key doesn't exist in the map yet, add it in
+    ///         // the obvious way.
+    ///         |_k, v| vec![v]);
+    /// }
+    ///
+    /// assert_eq!(map.len(), 3);
+    /// assert_eq!(map.get(&"a key"), &vec!["value", "new value"]);
+    /// assert_eq!(map.get(&"b key"), &vec!["new value"]);
+    /// assert_eq!(map.get(&"z key"), &vec!["new value", "value"]);
+    /// ```
+    pub fn find_with_or_insert_with<'a, A>(&'a mut self,
+                                           k: K,
+                                           a: A,
+                                           found: |&K, &mut V, A|,
+                                           not_found: |&K, A| -> V)
+                                          -> &'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
             }
         }