From a343e435d59e21188f6fc5918324c751a7eff6a9 Mon Sep 17 00:00:00 2001
From: Kevin Cantu <me@kevincantu.org>
Date: Sun, 11 Nov 2012 01:01:37 -0800
Subject: [PATCH] Add an insert_with_key function to the Map trait

---
 src/libstd/map.rs                             | 25 ++++++++++++++++++-
 src/libstd/smallintmap.rs                     |  7 ++++++
 .../class-impl-very-parameterized-trait.rs    |  8 ++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

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))
     }
diff --git a/src/test/run-pass/class-impl-very-parameterized-trait.rs b/src/test/run-pass/class-impl-very-parameterized-trait.rs
index c11e41eb57f..c8a08dbf818 100644
--- a/src/test/run-pass/class-impl-very-parameterized-trait.rs
+++ b/src/test/run-pass/class-impl-very-parameterized-trait.rs
@@ -61,6 +61,14 @@ impl<T: Copy> cat<T> : Map<int, T> {
      else { None }
   }
 
+  fn insert_with_key(ff: fn(+k: int, +v0: T, +v1: T) -> T, +key: int, +val: T) -> bool {
+    match self.find(key) {
+      None            => return self.insert(key, val),
+      Some(copy orig) => return self.insert(key, ff(key, orig, val))
+    }
+  }
+
+
   fn remove(+k:int) -> bool {
     match self.find(k) {
       Some(x) => {