auto merge of #5365 : thestinger/rust/map, r=catamorphism

2 files changed, 30 insertions, 1 deletions

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 56ee3bd5893..8d4ff994195 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) {