Add an insert_with_key function to the Map trait

2 files changed, 31 insertions, 1 deletions

diff --git a/src/libstd/map.rs b/src/libstd/map.rs
index 915202143a1..d6feebc016c 100644
--- a/src/libstd/map.rs
+++ b/src/libstd/map.rs
@@ -27,7 +27,15 @@ pub trait Map<K:Eq IterBytes Hash Copy, V: Copy> {
      *
      * Returns true if the key did not already exist in the map
      */
-    fn insert(v: K, v: V) -> bool;
+    fn insert(key: K, value: V) -> bool;
+
+    /**
+     * Add a value to the map.
+     *
+     * If the map contains a value for the key, use the function
+     * to set a new value.
+     */
+    fn insert_with_key(ff: fn(K, V, V) -> V, key: K, value: V) -> bool;
 
     /// Returns true if the map contains a value for the specified key
     pure fn contains_key(key: K) -> bool;
@@ -264,6 +272,14 @@ pub mod chained {
             }
         }
 
+        fn insert_with_key(ff: fn(K, V, V) -> V, key: K, val: V) -> bool {
+            // this can be optimized but first lets see if it compiles...
+            match self.find(key) {
+                None            => return self.insert(key, val),
+                Some(copy orig) => return self.insert(key, ff(key, orig, val))
+            }
+        }
+
         pure fn get(k: K) -> V {
             let opt_v = self.find(k);
             if opt_v.is_none() {
@@ -447,6 +463,13 @@ impl<K: Eq IterBytes Hash Copy, V: Copy> @Mut<LinearMap<K, V>>:
         }
     }
 
+     fn insert_with_key(ff: fn(K, V, V) -> V, key: K, val: V) -> bool {
+         match self.find(key) {
+             None            => return self.insert(key, val),
+             Some(copy orig) => return self.insert(key, ff(key, orig, val)),
+         }
+     }
+
     fn remove(key: K) -> bool {
         do self.borrow_mut |p| {
             p.remove(&key)
diff --git a/src/libstd/smallintmap.rs b/src/libstd/smallintmap.rs
index 9dc216a2155..8439d214ca0 100644
--- a/src/libstd/smallintmap.rs
+++ b/src/libstd/smallintmap.rs
@@ -103,6 +103,13 @@ impl<V: Copy> SmallIntMap<V>: map::Map<uint, V> {
     pure fn find(key: uint) -> Option<V> { find(self, key) }
     fn rehash() { fail }
 
+    fn insert_with_key(ff: fn(uint, V, V) -> V, key: uint, val: V) -> bool {
+        match self.find(key) {
+            None            => return self.insert(key, val),
+            Some(copy orig) => return self.insert(key, ff(key, orig, val)),
+        }
+    }
+
     pure fn each(it: fn(key: uint, value: V) -> bool) {
         self.each_ref(|k, v| it(*k, *v))
     }