Add pop() and swap() to the Map trait

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 {
-        if *key >= self.v.len() {
-            return false;
+        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 => ()
         }
-        let removed = self.v[*key].is_some();
-        self.v[*key] = None;
-        removed
+        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 None;
+        }
+        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);
+    }
 }