add the mutate_values method to the Map trait

6 files changed, 61 insertions, 12 deletions

diff --git a/src/libcore/container.rs b/src/libcore/container.rs
index 4f1f6004aad..828678b65f7 100644
--- a/src/libcore/container.rs
+++ b/src/libcore/container.rs
@@ -35,6 +35,9 @@ pub trait Map<K, V>: Mutable {
     /// Visit all values
     pure fn each_value(&self, f: &fn(&V) -> bool);
 
+    /// Iterate over the map and mutate the contained values
+    fn mutate_values(&mut self, f: &fn(&K, &mut V) -> bool);
+
     /// Return the value corresponding to the key in the map
     pure fn find(&self, key: &K) -> Option<&self/V>;
 
diff --git a/src/libcore/hashmap.rs b/src/libcore/hashmap.rs
index 2adcee495a7..10d188a8555 100644
--- a/src/libcore/hashmap.rs
+++ b/src/libcore/hashmap.rs
@@ -324,6 +324,19 @@ pub mod linear {
             self.each(|&(_, v)| blk(v))
         }
 
+        /// Iterate over the map and mutate the contained values
+        fn mutate_values(&mut self, blk: &fn(&'self K,
+                              &'self mut V) -> bool) {
+            for uint::range(0, self.buckets.len()) |i| {
+                match self.buckets[i] {
+                  Some(Bucket{key: ref key, value: ref mut value, _}) => {
+                    if !blk(key, value) { return }
+                  }
+                  None => ()
+                }
+            }
+        }
+
         /// Return the value corresponding to the key in the map
         pure fn find(&self, k: &K) -> Option<&self/V> {
             match self.bucket_for_key(k) {
diff --git a/src/libcore/trie.rs b/src/libcore/trie.rs
index 7dc85cba297..966db4ec662 100644
--- a/src/libcore/trie.rs
+++ b/src/libcore/trie.rs
@@ -81,11 +81,16 @@ impl<T> Map<uint, T> for TrieMap<T> {
 
     /// Visit all values in order
     #[inline(always)]
-    pure fn each_value(&self,
-                       f: &fn(&T) -> bool) {
+    pure fn each_value(&self, f: &fn(&T) -> bool) {
         self.each(|&(_, v)| f(v))
     }
 
+    /// Iterate over the map and mutate the contained values
+    #[inline(always)]
+    fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) {
+        self.root.mutate_values(f);
+    }
+
     /// Return the value corresponding to the key in the map
     #[inline(hint)]
     pure fn find(&self, key: &uint) -> Option<&self/T> {
@@ -150,11 +155,6 @@ impl<T> TrieMap<T> {
     pure fn each_value_reverse(&self, f: &fn(&T) -> bool) {
         self.each_reverse(|&(_, v)| f(v))
     }
-
-    /// Iterate over the map and mutate the contained values
-    fn mutate_values(&mut self, f: &fn(uint, &mut T) -> bool) {
-        self.root.mutate_values(f);
-    }
 }
 
 pub struct TrieSet {
@@ -248,13 +248,13 @@ impl<T> TrieNode<T> {
         true
     }
 
-    fn mutate_values(&mut self, f: &fn(uint, &mut T) -> bool) -> bool {
+    fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) -> bool {
         for vec::each_mut(self.children) |child| {
             match *child {
                 Internal(ref mut x) => if !x.mutate_values(f) {
                     return false
                 },
-                External(k, ref mut v) => if !f(k, v) { return false },
+                External(k, ref mut v) => if !f(&k, v) { return false },
                 Nothing => ()
             }
         }
@@ -269,8 +269,8 @@ pure fn chunk(n: uint, idx: uint) -> uint {
     (n >> (SHIFT * real_idx)) & MASK
 }
 
-fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint,
-                   value: T, idx: uint) -> bool {
+fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
+             idx: uint) -> bool {
     let mut tmp = Nothing;
     tmp <-> *child;
     let mut added = false;
diff --git a/src/libstd/smallintmap.rs b/src/libstd/smallintmap.rs
index 726e7c36abd..bdce257e347 100644
--- a/src/libstd/smallintmap.rs
+++ b/src/libstd/smallintmap.rs
@@ -85,7 +85,17 @@ impl<V> Map<uint, V> for SmallIntMap<V> {
         self.each(|&(_, v)| blk(v))
     }
 
-    /// Return the value corresponding to the key in the map
+    /// Visit all key-value pairs in order
+    fn mutate_values(&mut self, it: &fn(&uint, &'self mut V) -> bool) {
+        for uint::range(0, self.v.len()) |i| {
+            match self.v[i] {
+              Some(ref mut elt) => if !it(&i, elt) { break },
+              None => ()
+            }
+        }
+    }
+
+    /// Iterate over the map and mutate the contained values
     pure fn find(&self, key: &uint) -> Option<&self/V> {
         if *key < self.v.len() {
             match self.v[*key] {
diff --git a/src/libstd/treemap.rs b/src/libstd/treemap.rs
index 42a84da43d2..0a4c980f76a 100644
--- a/src/libstd/treemap.rs
+++ b/src/libstd/treemap.rs
@@ -134,6 +134,11 @@ impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
         self.each(|&(_, v)| f(v))
     }
 
+    /// Iterate over the map and mutate the contained values
+    fn mutate_values(&mut self, f: &fn(&'self K, &'self mut V) -> bool) {
+        mutate_values(&mut self.root, f);
+    }
+
     /// Return the value corresponding to the key in the map
     pure fn find(&self, key: &K) -> Option<&self/V> {
         let mut current: &self/Option<~TreeNode<K, V>> = &self.root;
@@ -558,6 +563,20 @@ pure fn each_reverse<K: TotalOrd, V>(node: &r/Option<~TreeNode<K, V>>,
     }
 }
 
+fn mutate_values<K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
+                                 f: &fn(&'r K, &'r mut V) -> bool) -> bool {
+    match *node {
+      Some(~TreeNode{key: ref key, value: ref mut value, left: ref mut left,
+                     right: ref mut right, _}) => {
+        if !mutate_values(left, f) { return false }
+        if !f(key, value) { return false }
+        if !mutate_values(right, f) { return false }
+      }
+      None => return false
+    }
+    true
+}
+
 // Remove left horizontal link by rotating right
 fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
     if node.left.map_default(false, |x| x.level == node.level) {
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 6cb0749ddb5..0c285258f75 100644
--- a/src/test/run-pass/class-impl-very-parameterized-trait.rs
+++ b/src/test/run-pass/class-impl-very-parameterized-trait.rs
@@ -81,6 +81,10 @@ impl<T> Map<int, T> for cat<T> {
         for self.each |&(_, v)| { if !f(v) { break; } loop;};
     }
 
+    fn mutate_values(&mut self, f: &fn(&int, &mut T) -> bool) {
+        fail!(~"nope")
+    }
+
     fn insert(&mut self, k: int, _: T) -> bool {
         self.meows += k;
         true