Add member operations to collections

src/libcollections/btree/map.rs

@@ -213,10 +213,32 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V> where K: Borrow<Q>, Q: Ord {
+        self.keyed_get(key).map(|x| x.1)
+    }
+
+    /// Returns a reference to the key and the value corresponding to the key.
+    ///
+    /// The key may be any borrowed form of the map's key type, but the ordering
+    /// on the borrowed form *must* match the ordering on the key type.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map = BTreeMap::new();
+    /// map.insert(1, "a");
+    /// assert_eq!(map.keyed_get(&1), Some((&1, &"a")));
+    /// assert_eq!(map.get(&2), None);
+    /// ```
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_get<Q: ?Sized>(&self, key: &Q) -> Option<(&K, &V)> where K: Borrow<Q>, Q: Ord {
         let mut cur_node = &self.root;
         loop {
             match Node::search(cur_node, key) {
-                Found(handle) => return Some(handle.into_kv().1),
+                Found(handle) => return Some(handle.into_kv()),
                 GoDown(handle) => match handle.force() {
                     Leaf(_) => return None,
                     Internal(internal_handle) => {
@@ -268,12 +290,42 @@ impl<K: Ord, V> BTreeMap<K, V> {
     // See `get` for implementation notes, this is basically a copy-paste with mut's added
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V> where K: Borrow<Q>, Q: Ord {
+        self.keyed_get_mut(key).map(|x| x.1)
+    }
+
+    /// Returns a mutable reference to the value corresponding to the key.
+    ///
+    /// The key may be any borrowed form of the map's key type, but the ordering
+    /// on the borrowed form *must* match the ordering on the key type.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map = BTreeMap::new();
+    /// map.insert(1, "a");
+    /// if let Some((key, x)) = map.keyed_get_mut(&1) {
+    ///     assert_eq!(key, &1);
+    ///     *x = "b";
+    /// }
+    /// assert_eq!(map[&1], "b");
+    /// ```
+    // See `keyed_get` for implementation notes, this is basically a copy-paste with mut's added
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<(&K, &mut V)> where
+            K: Borrow<Q>, Q: Ord {
         // temp_node is a Borrowck hack for having a mutable value outlive a loop iteration
         let mut temp_node = &mut self.root;
         loop {
             let cur_node = temp_node;
             match Node::search(cur_node, key) {
-                Found(handle) => return Some(handle.into_kv_mut().1),
+                Found(handle) => {
+                    let (key, value) = handle.into_kv_mut();
+                    return Some((&*key, value));
+                },
                 GoDown(handle) => match handle.force() {
                     Leaf(_) => return None,
                     Internal(internal_handle) => {
@@ -328,7 +380,31 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// assert_eq!(map[&37], "c");
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
-    pub fn insert(&mut self, mut key: K, mut value: V) -> Option<V> {
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        self.keyed_insert(key, value).map(|x| x.1)
+    }
+
+    /// Inserts a key-value pair into the map. If the key already had a value
+    /// present in the map, that key and value are returned. Otherwise,
+    /// `None` is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map = BTreeMap::new();
+    /// assert_eq!(map.keyed_insert(37, "a"), None);
+    /// assert_eq!(map.is_empty(), false);
+    ///
+    /// map.insert(37, "b");
+    /// assert_eq!(map.keyed_insert(37, "c"), Some((37, "b")));
+    /// assert_eq!(map[&37], "c");
+    /// ```
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_insert(&mut self, mut key: K, mut value: V) -> Option<(K, V)> {
         // This is a stack of rawptrs to nodes paired with indices, respectively
         // representing the nodes and edges of our search path. We have to store rawptrs
         // because as far as Rust is concerned, we can mutate aliased data with such a
@@ -354,7 +430,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
                     Found(mut handle) => {
                         // Perfect match, swap the values and return the old one
                         mem::swap(handle.val_mut(), &mut value);
-                        Finished(Some(value))
+                        mem::swap(handle.key_mut(), &mut key);
+                        Finished(Some((key, value)))
                     },
                     GoDown(handle) => {
                         // We need to keep searching, try to get the search stack
@@ -438,14 +515,38 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V> where K: Borrow<Q>, Q: Ord {
+        self.keyed_remove(key).map(|x| x.1)
+    }
+
+    /// Removes a key from the map, returning the key and the value at the key
+    /// if the key was previously in the map.
+    ///
+    /// The key may be any borrowed form of the map's key type, but the ordering
+    /// on the borrowed form *must* match the ordering on the key type.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map = BTreeMap::new();
+    /// map.insert(1, "a");
+    /// assert_eq!(map.keyed_remove(&1), Some((1, "a")));
+    /// assert_eq!(map.keyed_remove(&1), None);
+    /// ```
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_remove<Q: ?Sized>(&mut self, key: &Q) -> Option<(K, V)> where
+            K: Borrow<Q>, Q: Ord {
         // See `swap` for a more thorough description of the stuff going on in here
         let mut stack = stack::PartialSearchStack::new(self);
         loop {
             let result = stack.with(move |pusher, node| {
                 match Node::search(node, key) {
                     Found(handle) => {
                         // Perfect match. Terminate the stack here, and remove the entry
-                        Finished(Some(pusher.seal(handle).remove()))
+                        Finished(Some(pusher.seal(handle).keyed_remove()))
                     },
                     GoDown(handle) => {
                         // We need to keep searching, try to go down the next edge
@@ -693,16 +794,16 @@ mod stack {
     impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::Leaf> {
         /// Removes the key and value in the top element of the stack, then handles underflows as
         /// described in BTree's pop function.
-        fn remove_leaf(mut self) -> V {
+        fn remove_leaf(mut self) -> (K, V) {
             self.map.length -= 1;
 
             // Remove the key-value pair from the leaf that this search stack points to.
             // Then, note if the leaf is underfull, and promptly forget the leaf and its ptr
             // to avoid ownership issues.
-            let (value, mut underflow) = unsafe {
-                let (_, value) = self.top.from_raw_mut().remove_as_leaf();
+            let (key, value, mut underflow) = unsafe {
+                let (key, value) = self.top.from_raw_mut().remove_as_leaf();
                 let underflow = self.top.from_raw().node().is_underfull();
-                (value, underflow)
+                (key, value, underflow)
             };
 
             loop {
@@ -717,7 +818,7 @@ mod stack {
                             self.map.depth -= 1;
                             self.map.root.hoist_lone_child();
                         }
-                        return value;
+                        return (key, value);
                     }
                     Some(mut handle) => {
                         if underflow {
@@ -728,7 +829,7 @@ mod stack {
                             }
                         } else {
                             // All done!
-                            return value;
+                            return (key, value);
                         }
                     }
                 }
@@ -740,6 +841,16 @@ mod stack {
         /// Removes the key and value in the top element of the stack, then handles underflows as
         /// described in BTree's pop function.
         pub fn remove(self) -> V {
+            // Ensure that the search stack goes to a leaf. This is necessary to perform deletion
+            // in a BTree. Note that this may put the tree in an inconsistent state (further
+            // described in into_leaf's comments), but this is immediately fixed by the
+            // removing the value we want to remove
+            self.into_leaf().remove_leaf().1
+        }
+
+        /// Removes the key and value in the top element of the stack, then handles underflows as
+        /// described in BTree's pop function.
+        pub fn keyed_remove(self) -> (K, V) {
             // Ensure that the search stack goes to a leaf. This is necessary to perform deletion
             // in a BTree. Note that this may put the tree in an inconsistent state (further
             // described in into_leaf's comments), but this is immediately fixed by the

src/libcollections/vec_map.rs

@@ -490,9 +490,29 @@ impl<V> VecMap<V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn get(&self, key: &usize) -> Option<&V> {
+        self.keyed_get(key).map(|x| x.1)
+    }
+
+    /// Returns a reference to the key and the value corresponding to the key.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(vecmap)]
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::VecMap;
+    ///
+    /// let mut map = VecMap::new();
+    /// map.insert(1, "a");
+    /// assert_eq!(map.keyed_get(&1), Some((&1, &"a")));
+    /// assert_eq!(map.keyed_get(&2), None);
+    /// ```
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_get<'a>(&self, key: &'a usize) -> Option<(&'a usize, &V)> {
         if *key < self.v.len() {
             match self.v[*key] {
-              Some(ref value) => Some(value),
+              Some(ref value) => Some((key, value)),
               None => None
             }
         } else {
@@ -536,9 +556,32 @@ impl<V> VecMap<V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn get_mut(&mut self, key: &usize) -> Option<&mut V> {
+        self.keyed_get_mut(key).map(|x| x.1)
+    }
+
+    /// Returns a mutable reference to the value corresponding to the key.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collection_keyed)]
+    /// # #![feature(vecmap)]
+    /// use std::collections::VecMap;
+    ///
+    /// let mut map = VecMap::new();
+    /// map.insert(1, "a");
+    /// if let Some((key, x)) = map.keyed_get_mut(&1) {
+    ///     assert_eq!(key, &1);
+    ///     *x = "b";
+    /// }
+    /// assert_eq!(map[1], "b");
+    /// ```
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_get_mut<'a>(&mut self, key: &'a usize) -> Option<(&'a usize, &mut V)> {
         if *key < self.v.len() {
             match *(&mut self.v[*key]) {
-              Some(ref mut value) => Some(value),
+              Some(ref mut value) => Some((key, value)),
               None => None
             }
         } else {
@@ -565,11 +608,36 @@ impl<V> VecMap<V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn insert(&mut self, key: usize, value: V) -> Option<V> {
+        self.keyed_insert(key, value).map(|x| x.1)
+    }
+
+    /// Inserts a key-value pair into the map. If the key already had a value
+    /// present in the map, that key and value are returned. Otherwise,
+    /// `None` is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(vecmap)]
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::VecMap;
+    ///
+    /// let mut map = VecMap::new();
+    /// assert_eq!(map.keyed_insert(37, "a"), None);
+    /// assert_eq!(map.is_empty(), false);
+    ///
+    /// map.insert(37, "b");
+    /// assert_eq!(map.keyed_insert(37, "c"), Some((37, "b")));
+    /// assert_eq!(map[37], "c");
+    /// ```
+    #[unstable(feature = "collection_keyed",
+            reason="keyed was recently added")]
+    pub fn keyed_insert(&mut self, key: usize, value: V) -> Option<(usize, V)> {
         let len = self.v.len();
         if len <= key {
             self.v.extend((0..key - len + 1).map(|_| None));
         }
-        replace(&mut self.v[key], Some(value))
+        replace(&mut self.v[key], Some(value)).map(|x| (key, x))
     }
 
     /// Removes a key from the map, returning the value at the key if the key
@@ -588,11 +656,31 @@ impl<V> VecMap<V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn remove(&mut self, key: &usize) -> Option<V> {
+        self.keyed_remove(key).map(|x| x.1)
+    }
+
+    /// Removes a key from the map, returning the key and value at the key if
+    /// it was previously in the map.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(vecmap)]
+    /// # #![feature(collection_keyed)]
+    /// use std::collections::VecMap;
+    ///
+    /// let mut map = VecMap::new();
+    /// map.insert(1, "a");
+    /// assert_eq!(map.keyed_remove(&1), Some((1, "a")));
+    /// assert_eq!(map.keyed_remove(&1), None);
+    /// ```
+    #[stable(feature = "rust1", since = "1.0.0")]
+    pub fn keyed_remove(&mut self, key: &usize) -> Option<(usize, V)> {
         if *key >= self.v.len() {
             return None;
         }
         let result = &mut self.v[*key];
-        result.take()
+        result.take().map(|x| (*key, x))
     }
 
     /// Gets the given key's corresponding entry in the map for in-place manipulation.