Add Map::insert_with_key

src/libstd/map.rs

@@ -27,7 +27,24 @@ 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(key: K, newval: V, ff: fn(K, V, V) -> 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. (Like insert_with_key, but with a function
+ * of only values.)
+ */
+ fn insert_with(key: K, newval: V, ff: fn(V, V) -> V) -> bool;
 
  /// Returns true if the map contains a value for the specified key
  pure fn contains_key(key: K) -> bool;
@@ -264,6 +281,59 @@ pub mod chained {
  }
  }
 
+ fn insert_with_key(key: K, newval: V, ff: fn(K, V, V) -> V) -> bool {
+/*
+ match self.find(key) {
+ None => return self.insert(key, val),
+ Some(copy orig) => return self.insert(key, ff(key, orig, val))
+ }
+*/
+
+ let hash = key.hash_keyed(0,0) as uint;
+ match self.search_tbl(&key, hash) {
+ NotFound => {
+ self.count += 1u;
+ let idx = hash % vec::len(self.chains);
+ let old_chain = self.chains[idx];
+ self.chains[idx] = Some(@Entry {
+ hash: hash,
+ key: key,
+ value: newval,
+ next: old_chain});
+
+ // consider rehashing if more 3/4 full
+ let nchains = vec::len(self.chains);
+ let load = {num: (self.count + 1u) as int,
+ den: nchains as int};
+ if !util::rational_leq(load, {num:3, den:4}) {
+ self.rehash();
+ }
+
+ return true;
+ }
+ FoundFirst(idx, entry) => {
+ self.chains[idx] = Some(@Entry {
+ hash: hash,
+ key: key,
+ value: ff(key, entry.value, newval),
+ next: entry.next});
+ return false;
+ }
+ FoundAfter(prev, entry) => {
+ prev.next = Some(@Entry {
+ hash: hash,
+ key: key,
+ value: ff(key, entry.value, newval),
+ next: entry.next});
+ return false;
+ }
+ }
+ }
+
+ fn insert_with(key: K, newval: V, ff: fn(V, V) -> V) -> bool {
+ return self.insert_with_key(key, newval, |_k, v, v1| ff(v,v1));
+ }
+
  pure fn get(k: K) -> V {
  let opt_v = self.find(k);
  if opt_v.is_none() {
@@ -447,6 +517,17 @@ impl<K: Eq IterBytes Hash Copy, V: Copy> @Mut<LinearMap<K, V>>:
  }
  }
 
+ fn insert_with_key(key: K, newval: V, ff: fn(K, V, V) -> V) -> bool {
+ match self.find(key) {
+ None => return self.insert(key, newval),
+ Some(copy orig) => return self.insert(key, ff(key, orig, newval))
+ }
+ }
+
+ fn insert_with(key: K, newval: V, ff: fn(V, V) -> V) -> bool {
+ return self.insert_with_key(key, newval, |_k, v, v1| ff(v,v1));
+ }
+
  fn remove(key: K) -> bool {
  do self.borrow_mut |p| {
  p.remove(&key)
@@ -750,4 +831,35 @@ mod tests {
  assert map.get(~"b") == 2;
  assert map.get(~"c") == 3;
  }
+
+ #[test]
+ fn test_insert_with_key() {
+ let map = map::HashMap::<~str, uint>();
+
+ // given a new key, initialize it with this new count, given
+ // given an existing key, add more to its count
+ fn addMoreToCount(_k: ~str, v0: uint, v1: uint) -> uint {
+ v0 + v1
+ }
+
+ fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
+ v0 + v1
+ }
+
+ // count the number of several types of animal,
+ // adding in groups as we go
+ map.insert_with(~"cat", 1, addMoreToCount_simple);
+ map.insert_with_key(~"mongoose", 1, addMoreToCount);
+ map.insert_with(~"cat", 7, addMoreToCount_simple);
+ map.insert_with_key(~"ferret", 3, addMoreToCount);
+ map.insert_with_key(~"cat", 2, addMoreToCount);
+
+ // check the total counts
+ assert 10 == option::get(map.find(~"cat"));
+ assert 3 == option::get(map.find(~"ferret"));
+ assert 1 == option::get(map.find(~"mongoose"));
+
+ // sadly, no mythical animals were counted!
+ assert None == map.find(~"unicorn");
+ }
 }

src/libstd/smallintmap.rs

@@ -103,6 +103,17 @@ 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(key: uint, val: V, ff: fn(uint, V, V) -> V) -> bool {
+ match self.find(key) {
+ None => return self.insert(key, val),
+ Some(copy orig) => return self.insert(key, ff(key, orig, val)),
+ }
+ }
+
+ fn insert_with(key: uint, newval: V, ff: fn(V, V) -> V) -> bool {
+ return self.insert_with_key(key, newval, |_k, v, v1| ff(v,v1));
+ }
+
  pure fn each(it: fn(key: uint, value: V) -> bool) {
  self.each_ref(|k, v| it(*k, *v))
  }
@@ -142,3 +153,37 @@ impl<V: Copy> SmallIntMap<V>: ops::Index<uint, V> {
 pub fn as_map<V: Copy>(s: SmallIntMap<V>) -> map::Map<uint, V> {
  s as map::Map::<uint, V>
 }
+
+#[cfg(test)]
+mod tests {
+
+ #[test]
+ fn test_insert_with_key() {
+ let map: SmallIntMap<uint> = mk();
+
+ // given a new key, initialize it with this new count, given
+ // given an existing key, add more to its count
+ fn addMoreToCount(_k: uint, v0: uint, v1: uint) -> uint {
+ v0 + v1
+ }
+
+ fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
+ v0 + v1
+ }
+
+ // count integers
+ map.insert_with(3, 1, addMoreToCount_simple);
+ map.insert_with_key(9, 1, addMoreToCount);
+ map.insert_with(3, 7, addMoreToCount_simple);
+ map.insert_with_key(5, 3, addMoreToCount);
+ map.insert_with_key(3, 2, addMoreToCount);
+
+ // check the total counts
+ assert 10 == option::get(map.find(3));
+ assert 3 == option::get(map.find(5));
+ assert 1 == option::get(map.find(9));
+
+ // sadly, no sevens were counted
+ assert None == map.find(7);
+ }
+}

src/test/bench/shootout-k-nucleotide-pipes.rs

@@ -69,10 +69,7 @@ fn find(mm: HashMap<~[u8], uint>, key: ~str) -> uint {
 // given a map, increment the counter for a key
 fn update_freq(mm: HashMap<~[u8], uint>, key: &[u8]) {
  let key = vec::slice(key, 0, key.len());
- match mm.find(key) {
- option::None => { mm.insert(key, 1u ); }
- option::Some(val) => { mm.insert(key, 1u + val); }
- }
+ mm.insert_with(key, 1, |v,v1| { v+v1 });
 }
 
 // given a ~[u8], for each window call a function

src/test/bench/shootout-k-nucleotide.rs

@@ -66,10 +66,7 @@ fn find(mm: HashMap<~[u8], uint>, key: ~str) -> uint {
 // given a map, increment the counter for a key
 fn update_freq(mm: HashMap<~[u8], uint>, key: &[u8]) {
  let key = vec::slice(key, 0, key.len());
- match mm.find(key) {
- option::None => { mm.insert(key, 1u ); }
- option::Some(val) => { mm.insert(key, 1u + val); }
- }
+ mm.insert_with(key, 1, |v,v1| { v+v1 });
 }
 
 // given a ~[u8], for each window call a function

src/test/run-pass/class-impl-very-parameterized-trait.rs

@@ -61,6 +61,21 @@ impl<T: Copy> cat<T> : Map<int, T> {
  else { None }
  }
 
+ fn insert_with_key(+key: int, +val: T, ff: fn(+k: int, +v0: T, +v1: T) -> T) -> bool {
+ match self.find(key) {
+ None => return self.insert(key, val),
+ Some(copy orig) => return self.insert(key, ff(key, orig, val))
+ }
+ }
+
+ fn insert_with(+key: int, +val: T, ff: fn(+v0: T, +v1: T) -> T) -> bool {
+ match self.find(key) {
+ None => return self.insert(key, val),
+ Some(copy orig) => return self.insert(key, ff(orig, val))
+ }
+ }
+
+
  fn remove(+k:int) -> bool {
  match self.find(k) {
  Some(x) => {