Map trait improvements

src/libextra/smallintmap.rs

@@ -23,6 +23,7 @@ use std::vec;
 #[allow(missing_doc)]
 pub struct SmallIntMap<T> {
     priv v: ~[Option<T>],
+    priv last_key: uint,
 }
 
 impl<V> Container for SmallIntMap<V> {
@@ -69,17 +70,19 @@ impl<V> MutableMap<uint, V> for SmallIntMap<V> {
         }
     }
 
-    /// Insert a key-value pair into the map. An existing value for a
-    /// key is replaced by the new value. Return true if the key did
-    /// not already exist in the map.
-    fn insert(&mut self, key: uint, value: V) -> bool {
-        let exists = self.contains_key(&key);
+    /// Return the value corresponding to the key in the map, or create,
+    /// insert, and return a new value if it doesn't exist.
+    fn find_or_insert_with<'a>(&'a mut self, key: uint, f: &fn(&uint) -> V)
+                               -> (&'a uint, &'a mut V) {
         let len = self.v.len();
         if len <= key {
             self.v.grow_fn(key - len + 1, |_| None);
         }
-        self.v[key] = Some(value);
-        !exists
+        self.last_key = key;
+        if self.v[key].is_none() {
+            self.v[key] = Some(f(&key));
+        }
+        (&self.last_key, self.v[key].get_mut_ref())
     }
 
     /// Remove a key-value pair from the map. Return true if the key
@@ -88,17 +91,6 @@ impl<V> MutableMap<uint, V> for SmallIntMap<V> {
         self.pop(key).is_some()
     }
 
-    /// Insert a key-value pair from the map. If the key already had a value
-    /// present in the map, that value is returned. Otherwise None is returned.
-    fn swap(&mut self, key: uint, value: V) -> Option<V> {
-        match self.find_mut(&key) {
-            Some(loc) => { return Some(replace(loc, value)); }
-            None => ()
-        }
-        self.insert(key, value);
-        return None;
-    }
-
     /// Removes a key from the map, returning the value at the key if the key
     /// was previously in the map.
     fn pop(&mut self, key: &uint) -> Option<V> {
@@ -107,11 +99,15 @@ impl<V> MutableMap<uint, V> for SmallIntMap<V> {
         }
         self.v[*key].take()
     }
+
+    /// Does nothing for this implementation.
+    fn reserve_at_least(&mut self, _: uint) {
+    }
 }
 
 impl<V> SmallIntMap<V> {
     /// Create an empty SmallIntMap
-    pub fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[]} }
+    pub fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[], last_key: 0} }
 
     pub fn get<'a>(&'a self, key: &uint) -> &'a V {
         self.find(key).expect("key not present")
@@ -337,6 +333,39 @@ mod test_map {
         assert_eq!(m.swap(1, 4), Some(3));
     }
 
+    #[test]
+    fn test_find_or_insert() {
+        let mut m = SmallIntMap::new();
+        {
+            let (k, v) = m.find_or_insert(1, 2);
+            assert_eq!((*k, *v), (1, 2));
+        }
+        {
+            let (k, v) = m.find_or_insert(1, 3);
+            assert_eq!((*k, *v), (1, 2));
+        }
+    }
+
+    #[test]
+    fn test_find_or_insert_with() {
+        let mut m = SmallIntMap::new();
+        {
+            let (k, v) = m.find_or_insert_with(1, |_| 2);
+            assert_eq!((*k, *v), (1, 2));
+        }
+        {
+            let (k, v) = m.find_or_insert_with(1, |_| 3);
+            assert_eq!((*k, *v), (1, 2));
+        }
+    }
+
+    #[test]
+    fn test_insert_or_update_with() {
+        let mut m = SmallIntMap::new();
+        assert_eq!(*m.insert_or_update_with(1, 2, |_,x| *x+=1), 2);
+        assert_eq!(*m.insert_or_update_with(1, 2, |_,x| *x+=1), 3);
+    }
+
     #[test]
     fn test_pop() {
         let mut m = SmallIntMap::new();

src/libextra/treemap.rs

@@ -16,6 +16,7 @@
 use std::util::{swap, replace};
 use std::iter::{Peekable};
 use std::cmp::Ordering;
+use std::cell::Cell;
 
 // This is implemented as an AA tree, which is a simplified variation of
 // a red-black tree where red (horizontal) nodes can only be added
@@ -116,19 +117,52 @@ impl<K: TotalOrd, V> MutableMap<K, V> for TreeMap<K, V> {
 
     /// Insert a key-value pair from the map. If the key already had a value
     /// present in the map, that value is returned. Otherwise None is returned.
-    fn swap(&mut self, key: K, value: V) -> Option<V> {
-        let ret = insert(&mut self.root, key, value);
-        if ret.is_none() { self.length += 1 }
+    fn swap(&mut self, k: K, v: V) -> Option<V> {
+        let cell = Cell::new(v);
+        let mut ret = None;
+        insert_helper(&mut self.root,
+                      k,
+                      |_| { self.length += 1; cell.take() },
+                      |old| { ret = Some(replace(old, cell.take()))},
+                      |_| {});
         ret
     }
 
+
+    /// Return the value corresponding to the key in the map, or create,
+    /// insert, and return a new value if it doesn't exist.
+    fn find_or_insert_with<'a>(&'a mut self, k: K, f: &fn(&K) -> V)
+                               -> (&'a K, &'a mut V) {
+        // since insert_helper needs to rebalance the tree on the
+        // way back up the stack, it can't safely hold a reference
+        // to the found or newly inserted node.  Instead we will
+        // build a trail back to the node we want to find.
+        let mut trail = ~[];
+        insert_helper(&mut self.root, k,
+                      |k| { self.length += 1; f(k) },
+                      |_| {},
+                      |o| { trail.push(o) });
+        let node = trail.rev_iter().fold(self.root.get_mut_ref(), |node, dir| {
+            match *dir {
+                Less => &mut node.left,
+                Greater => &mut node.right,
+                Equal => unreachable!(),
+            }.get_mut_ref()
+        });
+        (&node.key, &mut node.value)
+    }
+
     /// Removes a key from the map, returning the value at the key if the key
     /// was previously in the map.
     fn pop(&mut self, key: &K) -> Option<V> {
         let ret = remove(&mut self.root, key);
         if ret.is_some() { self.length -= 1 }
         ret
     }
+
+    /// Does nothing for this implementation.
+    fn reserve_at_least(&mut self, _: uint) {
+    }
 }
 
 impl<K: TotalOrd, V> TreeMap<K, V> {
@@ -694,25 +728,31 @@ fn mutate_values<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
 }
 
 // Remove left horizontal link by rotating right
-fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
+fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) -> bool {
     if node.left.as_ref().map_default(false, |x| x.level == node.level) {
         let mut save = node.left.take_unwrap();
         swap(&mut node.left, &mut save.right); // save.right now None
         swap(node, &mut save);
         node.right = Some(save);
+        true
+    } else {
+        false
     }
 }
 
 // Remove dual horizontal link by rotating left and increasing level of
 // the parent
-fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
+fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) -> bool {
     if node.right.as_ref().map_default(false,
       |x| x.right.as_ref().map_default(false, |y| y.level == node.level)) {
         let mut save = node.right.take_unwrap();
         swap(&mut node.right, &mut save.left); // save.left now None
         save.level += 1;
         swap(node, &mut save);
         node.left = Some(save);
+        true
+    } else {
+        false
     }
 }
 
@@ -731,33 +771,43 @@ fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
     }
 }
 
-fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
-                          key: K, value: V) -> Option<V> {
+fn insert_helper<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
+                                 key: K,
+                                 insert: &fn(&K) -> V,
+                                 update: &fn(&mut V),
+                                 build_trail: &fn(Ordering)) {
     match *node {
-      Some(ref mut save) => {
-        match key.cmp(&save.key) {
-          Less => {
-            let inserted = insert(&mut save.left, key, value);
-            skew(save);
-            split(save);
-            inserted
-          }
-          Greater => {
-            let inserted = insert(&mut save.right, key, value);
-            skew(save);
-            split(save);
-            inserted
-          }
-          Equal => {
-            save.key = key;
-            Some(replace(&mut save.value, value))
-          }
+        Some(ref mut save) => {
+            match key.cmp(&save.key) {
+                Less => {
+                    insert_helper(&mut save.left, key, insert, update,
+                                  |o| build_trail(o));
+                    if !skew(save) {
+                        build_trail(Less);
+                    }
+                    if split(save) {
+                        build_trail(Less);
+                    }
+                }
+                Greater => {
+                    insert_helper(&mut save.right, key, insert, update,
+                                  |o| build_trail(o));
+                    if skew(save) {
+                        build_trail(Greater);
+                    }
+                    if !split(save) {
+                        build_trail(Greater);
+                    }
+                }
+                Equal => {
+                    update(&mut save.value);
+                }
+            }
+        }
+        None => {
+            let value = insert(&key);
+            *node = Some(~TreeNode::new(key, value));
         }
-      }
-      None => {
-       *node = Some(~TreeNode::new(key, value));
-        None
-      }
     }
 }
 
@@ -828,11 +878,11 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
 
                 for right in save.right.mut_iter() {
                     skew(right);
-                    for x in right.right.mut_iter() { skew(x) }
+                    for x in right.right.mut_iter() { skew(x); }
                 }
 
                 split(save);
-                for x in save.right.mut_iter() { split(x) }
+                for x in save.right.mut_iter() { split(x); }
             }
 
             return ret;
@@ -923,6 +973,47 @@ mod test_treemap {
         assert_eq!(m.find(&2).unwrap(), &11);
     }
 
+    #[test]
+    fn test_find_or_insert_with() {
+        let mut m = TreeMap::new();
+        {
+            let (k, v) = m.find_or_insert_with(1, |_| 2);
+            assert_eq!((*k, *v), (1, 2));
+        }
+        {
+            let (k, v) = m.find_or_insert_with(1, |_| 3);
+            assert_eq!((*k, *v), (1, 2));
+        }
+    }
+
+    #[test]
+    fn test_find_or_insert() {
+        let mut m = TreeMap::new();
+        for i in range(0, 10) {
+            let (k, v) = m.find_or_insert(i, i + 1);
+            assert_eq!((*k, *v), (i, i + 1))
+        }
+        for i in range(0, 10) {
+            let (k, v) = m.find_or_insert(-i, i + 1);
+            assert_eq!((*k, *v), (-i, i + 1))
+        }
+        for i in range(0, 10) {
+            let (k, v) = m.find_or_insert(i, i + 2);
+            assert_eq!((*k, *v), (i, i + 1))
+        }
+        for i in range(0, 10) {
+            let (k, v) = m.find_or_insert(-i, i + 2);
+            assert_eq!((*k, *v), (-i, i + 1))
+        }
+    }
+
+    #[test]
+    fn test_insert_or_update_with() {
+        let mut m = TreeMap::new();
+        assert_eq!(*m.insert_or_update_with(1, 2, |_,x| *x+=1), 2);
+        assert_eq!(*m.insert_or_update_with(1, 2, |_,x| *x+=1), 3);
+    }
+
     #[test]
     fn test_clear() {
         let mut m = TreeMap::new();

src/librustc/middle/dataflow.rs

@@ -208,7 +208,7 @@ impl<O:DataFlowOperator> DataFlowContext<O> {
     fn compute_id_range(&mut self, id: ast::NodeId) -> (uint, uint) {
         let mut expanded = false;
         let len = self.nodeid_to_bitset.len();
-        let n = do self.nodeid_to_bitset.find_or_insert_with(id) |_| {
+        let (_, n) = do self.nodeid_to_bitset.find_or_insert_with(id) |_| {
             expanded = true;
             len
         };

src/librustc/middle/typeck/check/regionmanip.rs

@@ -36,7 +36,8 @@ pub fn replace_bound_regions_in_fn_sig(
                 debug!("region r={}", r.to_str());
                 match r {
                 ty::ReLateBound(s, br) if s == fn_sig.binder_id => {
-                    *map.find_or_insert_with(br, |_| mapf(br))
+                    let (_, v) = map.find_or_insert_with(br, |_| mapf(br));
+                    *v
                 }
                 _ => r
             }});

src/librustdoc/html/render.rs

@@ -491,7 +491,7 @@ impl DocFolder for Cache {
             clean::ImplItem(ref i) => {
                 match i.trait_ {
                     Some(clean::ResolvedPath{ id, _ }) => {
-                        let v = do self.implementors.find_or_insert_with(id) |_|{
+                        let (_, v) = do self.implementors.find_or_insert_with(id) |_|{
                             ~[]
                         };
                         match i.for_ {
@@ -594,7 +594,7 @@ impl DocFolder for Cache {
                     clean::Item{ attrs, inner: clean::ImplItem(i), _ } => {
                         match i.for_ {
                             clean::ResolvedPath { id, _ } => {
-                                let v = do self.impls.find_or_insert_with(id) |_| {
+                                let (_, v) = do self.impls.find_or_insert_with(id) |_| {
                                     ~[]
                                 };
                                 // extract relevant documentation for this impl
@@ -1607,7 +1607,7 @@ fn build_sidebar(m: &clean::Module) -> HashMap<~str, ~[~str]> {
             None => continue,
             Some(ref s) => s.to_owned(),
         };
-        let v = map.find_or_insert_with(short.to_owned(), |_| ~[]);
+        let (_, v) = map.find_or_insert_with(short.to_owned(), |_| ~[]);
         v.push(myname);
     }