auto merge of #6236 : alexcrichton/rust/more-map-methods, r=thestinger

Closes #5392 and #5393

I implemented the pop/swap methods for TrieMap/TreeMap/SmallIntMap, and I also updated all of them such that pop isn't just a remove/insert, but rather it's all done in one operation.

One thing I did notice is that with default methods it'd be really nice to define `insert` and `remove` in terms of `pop` and `swap` (or vice versa, just to have them available).

6 files changed, 187 insertions, 83 deletions

diff --git a/src/libcore/container.rs b/src/libcore/container.rs
index 00ea4a93221..37b904bbe63 100644
--- a/src/libcore/container.rs
+++ b/src/libcore/container.rs
@@ -55,6 +55,14 @@ pub trait Map<K, V>: Mutable {
     /// Remove a key-value pair from the map. Return true if the key
     /// was present in the map, otherwise false.
     fn remove(&mut self, key: &K) -> bool;
+
+    /// 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, k: K, v: V) -> Option<V>;
+
+    /// Removes a key from the map, returning the value at the key if the key
+    /// was previously in the map.
+    fn pop(&mut self, k: &K) -> Option<V>;
 }
 
 pub trait Set<T>: Mutable {
diff --git a/src/libcore/hashmap.rs b/src/libcore/hashmap.rs
index 9b01c1dad06..33fdc98137e 100644
--- a/src/libcore/hashmap.rs
+++ b/src/libcore/hashmap.rs
@@ -24,8 +24,8 @@ use rand::RngUtil;
 use rand;
 use uint;
 use vec;
-use util::unreachable;
 use kinds::Copy;
+use util::{replace, unreachable};
 
 static INITIAL_CAPACITY: uint = 32u; // 2^5
 
@@ -204,7 +204,7 @@ priv impl<K:Hash + Eq,V> HashMap<K, V> {
     /// Inserts the key value pair into the buckets.
     /// Assumes that there will be a bucket.
     /// True if there was no previous entry with that key
-    fn insert_internal(&mut self, hash: uint, k: K, v: V) -> bool {
+    fn insert_internal(&mut self, hash: uint, k: K, v: V) -> Option<V> {
         match self.bucket_for_key_with_hash(hash, &k) {
             TableFull => { fail!(~"Internal logic error"); }
             FoundHole(idx) => {
@@ -213,14 +213,19 @@ priv impl<K:Hash + Eq,V> HashMap<K, V> {
                 self.buckets[idx] = Some(Bucket{hash: hash, key: k,
                                                 value: v});
                 self.size += 1;
-                true
+                None
             }
             FoundEntry(idx) => {
                 debug!("insert overwrite (%?->%?) at idx %?, hash %?",
                        k, v, idx, hash);
-                self.buckets[idx] = Some(Bucket{hash: hash, key: k,
-                                                value: v});
-                false
+                match self.buckets[idx] {
+                    None => { fail!(~"insert_internal: Internal logic error") }
+                    Some(ref mut b) => {
+                        b.hash = hash;
+                        b.key = k;
+                        Some(replace(&mut b.value, v))
+                    }
+                }
             }
         }
     }
@@ -361,6 +366,20 @@ impl<K:Hash + Eq,V> Map<K, V> for HashMap<K, V> {
     /// key is replaced by the new value. Return true if the key did
     /// not already exist in the map.
     fn insert(&mut self, k: K, v: V) -> bool {
+        self.swap(k, v).is_none()
+    }
+
+    /// Remove a key-value pair from the map. Return true if the key
+    /// was present in the map, otherwise false.
+    fn remove(&mut self, k: &K) -> bool {
+        self.pop(k).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, k: K, v: V) -> Option<V> {
+        // this could be faster.
+
         if self.size >= self.resize_at {
             // n.b.: We could also do this after searching, so
             // that we do not resize if this call to insert is
@@ -375,10 +394,11 @@ impl<K:Hash + Eq,V> Map<K, V> for HashMap<K, V> {
         self.insert_internal(hash, k, v)
     }
 
-    /// Remove a key-value pair from the map. Return true if the key
-    /// was present in the map, otherwise false.
-    fn remove(&mut self, k: &K) -> bool {
-        self.pop(k).is_some()
+    /// Removes a key from the map, returning the value at the key if the key
+    /// was previously in the map.
+    fn pop(&mut self, k: &K) -> Option<V> {
+        let hash = k.hash_keyed(self.k0, self.k1) as uint;
+        self.pop_internal(hash, k)
     }
 }
 
@@ -402,31 +422,6 @@ pub impl<K: Hash + Eq, V> HashMap<K, V> {
         }
     }
 
-    fn pop(&mut self, k: &K) -> Option<V> {
-        let hash = k.hash_keyed(self.k0, self.k1) as uint;
-        self.pop_internal(hash, k)
-    }
-
-    fn swap(&mut self, k: K, v: V) -> Option<V> {
-        // this could be faster.
-        let hash = k.hash_keyed(self.k0, self.k1) as uint;
-        let old_value = self.pop_internal(hash, &k);
-
-        if self.size >= self.resize_at {
-            // n.b.: We could also do this after searching, so
-            // that we do not resize if this call to insert is
-            // simply going to update a key in place.  My sense
-            // though is that it's worse to have to search through
-            // buckets to find the right spot twice than to just
-            // resize in this corner case.
-            self.expand();
-        }
-
-        self.insert_internal(hash, k, v);
-
-        old_value
-    }
-
     /// Return the value corresponding to the key in the map, or insert
     /// and return the value if it doesn't exist.
     fn find_or_insert<'a>(&'a mut self, k: K, v: V) -> &'a V {
diff --git a/src/libcore/trie.rs b/src/libcore/trie.rs
index f0756be9944..380ad11ae92 100644
--- a/src/libcore/trie.rs
+++ b/src/libcore/trie.rs
@@ -11,6 +11,7 @@
 //! An ordered map and set for integer keys implemented as a radix trie
 
 use prelude::*;
+use util::{swap, replace};
 
 // FIXME: #5244: need to manually update the TrieNode constructor
 static SHIFT: uint = 4;
@@ -110,21 +111,33 @@ impl<T> Map<uint, T> for TrieMap<T> {
     /// not already exist in the map.
     #[inline(always)]
     fn insert(&mut self, key: uint, value: T) -> bool {
-        let ret = insert(&mut self.root.count,
-                         &mut self.root.children[chunk(key, 0)],
-                         key, value, 1);
-        if ret { self.length += 1 }
-        ret
+        self.swap(key, value).is_none()
     }
 
     /// Remove a key-value pair from the map. Return true if the key
     /// was present in the map, otherwise false.
     #[inline(always)]
     fn remove(&mut self, key: &uint) -> bool {
+        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: T) -> Option<T> {
+        let ret = insert(&mut self.root.count,
+                         &mut self.root.children[chunk(key, 0)],
+                         key, value, 1);
+        if ret.is_none() { self.length += 1 }
+        ret
+    }
+
+    /// 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<T> {
         let ret = remove(&mut self.root.count,
                          &mut self.root.children[chunk(*key, 0)],
                          *key, 1);
-        if ret { self.length -= 1 }
+        if ret.is_some() { self.length -= 1 }
         ret
     }
 }
@@ -289,14 +302,15 @@ fn find_mut<'r, T>(child: &'r mut Child<T>, key: uint, idx: uint)
 }
 
 fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
-             idx: uint) -> bool {
+             idx: uint) -> Option<T> {
     let mut tmp = Nothing;
-    tmp <-> *child;
-    let mut added = false;
+    let ret;
+    swap(&mut tmp, child);
 
     *child = match tmp {
       External(stored_key, stored_value) => {
           if stored_key == key {
+              ret = Some(stored_value);
               External(stored_key, value)
           } else {
               // conflict - split the node
@@ -304,46 +318,49 @@ fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
               insert(&mut new.count,
                      &mut new.children[chunk(stored_key, idx)],
                      stored_key, stored_value, idx + 1);
-              insert(&mut new.count, &mut new.children[chunk(key, idx)], key,
-                     value, idx + 1);
-              added = true;
+              ret = insert(&mut new.count, &mut new.children[chunk(key, idx)],
+                           key, value, idx + 1);
               Internal(new)
           }
       }
       Internal(x) => {
         let mut x = x;
-        added = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
-                       value, idx + 1);
+        ret = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
+                     value, idx + 1);
         Internal(x)
       }
       Nothing => {
         *count += 1;
-        added = true;
+        ret = None;
         External(key, value)
       }
     };
-    added
+    return ret;
 }
 
 fn remove<T>(count: &mut uint, child: &mut Child<T>, key: uint,
-             idx: uint) -> bool {
+             idx: uint) -> Option<T> {
     let (ret, this) = match *child {
-      External(stored, _) => {
-          if stored == key { (true, true) } else { (false, false) }
+      External(stored, _) if stored == key => {
+        match replace(child, Nothing) {
+            External(_, value) => (Some(value), true),
+            _ => fail!()
+        }
       }
+      External(*) => (None, false),
       Internal(ref mut x) => {
           let ret = remove(&mut x.count, &mut x.children[chunk(key, idx)],
                            key, idx + 1);
           (ret, x.count == 0)
       }
-      Nothing => (false, false)
+      Nothing => (None, false)
     };
 
     if this {
         *child = Nothing;
         *count -= 1;
     }
-    ret
+    return ret;
 }
 
 #[cfg(test)]
@@ -516,4 +533,20 @@ mod tests {
             i += 1;
         }
     }
+
+    #[test]
+    fn test_swap() {
+        let mut m = TrieMap::new();
+        assert!(m.swap(1, 2) == None);
+        assert!(m.swap(1, 3) == Some(2));
+        assert!(m.swap(1, 4) == Some(3));
+    }
+
+    #[test]
+    fn test_pop() {
+        let mut m = TrieMap::new();
+        m.insert(1, 2);
+        assert!(m.pop(&1) == Some(2));
+        assert!(m.pop(&1) == None);
+    }
 }
diff --git a/src/libstd/smallintmap.rs b/src/libstd/smallintmap.rs
index 1b72300a178..fc83a39cacf 100644
--- a/src/libstd/smallintmap.rs
+++ b/src/libstd/smallintmap.rs
@@ -16,6 +16,7 @@
 use core::container::{Container, Mutable, Map, Set};
 use core::old_iter::{BaseIter};
 use core::option::{Some, None};
+use core::util::replace;
 
 pub struct SmallIntMap<T> {
     priv v: ~[Option<T>],
@@ -119,12 +120,27 @@ impl<V> Map<uint, V> for SmallIntMap<V> {
     /// Remove a key-value pair from the map. Return true if the key
     /// was present in the map, otherwise false.
     fn remove(&mut self, key: &uint) -> bool {
+        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> {
         if *key >= self.v.len() {
-            return false;
+            return None;
         }
-        let removed = self.v[*key].is_some();
-        self.v[*key] = None;
-        removed
+        replace(&mut self.v[*key], None)
     }
 }
 
@@ -237,4 +253,20 @@ mod tests {
         // sadly, no sevens were counted
         assert!(map.find(&7).is_none());
     }
+
+    #[test]
+    fn test_swap() {
+        let mut m = SmallIntMap::new();
+        assert!(m.swap(1, 2) == None);
+        assert!(m.swap(1, 3) == Some(2));
+        assert!(m.swap(1, 4) == Some(3));
+    }
+
+    #[test]
+    fn test_pop() {
+        let mut m = SmallIntMap::new();
+        m.insert(1, 2);
+        assert!(m.pop(&1) == Some(2));
+        assert!(m.pop(&1) == None);
+    }
 }
diff --git a/src/libstd/treemap.rs b/src/libstd/treemap.rs
index 51695f2fa7d..c8ab48e65c0 100644
--- a/src/libstd/treemap.rs
+++ b/src/libstd/treemap.rs
@@ -13,6 +13,7 @@
 //! `TotalOrd`.
 
 use core::iterator::*;
+use core::util::replace;
 
 // This is implemented as an AA tree, which is a simplified variation of
 // a red-black tree where where red (horizontal) nodes can only be added
@@ -150,16 +151,28 @@ impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
     /// key is replaced by the new value. Return true if the key did
     /// not already exist in the map.
     fn insert(&mut self, key: K, value: V) -> bool {
-        let ret = insert(&mut self.root, key, value);
-        if ret { self.length += 1 }
-        ret
+        self.swap(key, value).is_none()
     }
 
     /// Remove a key-value pair from the map. Return true if the key
     /// was present in the map, otherwise false.
     fn remove(&mut self, key: &K) -> bool {
+        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: K, value: V) -> Option<V> {
+        let ret = insert(&mut self.root, key, value);
+        if ret.is_none() { self.length += 1 }
+        ret
+    }
+
+    /// 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 { self.length -= 1 }
+        if ret.is_some() { self.length -= 1 }
         ret
     }
 }
@@ -581,7 +594,8 @@ 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) -> bool {
+fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
+                          key: K, value: V) -> Option<V> {
     match *node {
       Some(ref mut save) => {
         match key.cmp(&save.key) {
@@ -599,20 +613,19 @@ fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K, value: V)
           }
           Equal => {
             save.key = key;
-            save.value = value;
-            false
+            Some(replace(&mut save.value, value))
           }
         }
       }
       None => {
        *node = Some(~TreeNode::new(key, value));
-        true
+        None
       }
     }
 }
 
 fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
-                          key: &K) -> bool {
+                          key: &K) -> Option<V> {
     fn heir_swap<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>,
                             child: &mut Option<~TreeNode<K, V>>) {
         // *could* be done without recursion, but it won't borrow check
@@ -628,12 +641,12 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
 
     match *node {
       None => {
-        return false // bottom of tree
+        return None; // bottom of tree
       }
       Some(ref mut save) => {
-        let (removed, this) = match key.cmp(&save.key) {
-          Less => (remove(&mut save.left, key), false),
-          Greater => (remove(&mut save.right, key), false),
+        let (ret, rebalance) = match key.cmp(&save.key) {
+          Less => (remove(&mut save.left, key), true),
+          Greater => (remove(&mut save.right, key), true),
           Equal => {
             if save.left.is_some() {
                 if save.right.is_some() {
@@ -645,21 +658,24 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
                         save.value <-> left.value;
                     }
                     save.left = Some(left);
-                    remove(&mut save.left, key);
+                    (remove(&mut save.left, key), true)
                 } else {
+                    let new = save.left.swap_unwrap();
+                    let ~TreeNode{value, _} = replace(save, new);
                     *save = save.left.swap_unwrap();
+                    (Some(value), true)
                 }
-                (true, false)
             } else if save.right.is_some() {
-                *save = save.right.swap_unwrap();
-                (true, false)
+                let new = save.right.swap_unwrap();
+                let ~TreeNode{value, _} = replace(save, new);
+                (Some(value), true)
             } else {
-                (true, true)
+                (None, false)
             }
           }
         };
 
-        if !this {
+        if rebalance {
             let left_level = save.left.map_default(0, |x| x.level);
             let right_level = save.right.map_default(0, |x| x.level);
 
@@ -682,13 +698,13 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
                 for save.right.each_mut |x| { split(x) }
             }
 
-            return removed;
+            return ret;
         }
       }
     }
-
-    *node = None;
-    true
+    return match replace(node, None) {
+        Some(~TreeNode{value, _}) => Some(value), None => fail!()
+    };
 }
 
 #[cfg(test)]
@@ -1217,4 +1233,20 @@ mod test_set {
         let result: Option<(&uint, & &'static str)> = z.next();
         assert!(result.is_none());
     }
+
+    #[test]
+    fn test_swap() {
+        let mut m = TreeMap::new();
+        assert!(m.swap(1, 2) == None);
+        assert!(m.swap(1, 3) == Some(2));
+        assert!(m.swap(1, 4) == Some(3));
+    }
+
+    #[test]
+    fn test_pop() {
+        let mut m = TreeMap::new();
+        m.insert(1, 2);
+        assert!(m.pop(&1) == Some(2));
+        assert!(m.pop(&1) == None);
+    }
 }
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 cf887758bff..b89bf064092 100644
--- a/src/test/run-pass/class-impl-very-parameterized-trait.rs
+++ b/src/test/run-pass/class-impl-very-parameterized-trait.rs
@@ -103,6 +103,10 @@ impl<T> Map<int, T> for cat<T> {
             false
         }
     }
+
+    fn pop(&mut self, _k: &int) -> Option<T> { fail!() }
+
+    fn swap(&mut self, _k: int, _v: T) -> Option<T> { fail!() }
 }
 
 pub impl<T> cat<T> {