implement RFC 1194

src/libcollections/btree/map.rs

@@ -459,7 +459,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
                 }
             });
             match result {
-                Finished(ret) => return ret,
+                Finished(ret) => return ret.map(|(_, v)| v),
                 Continue(new_stack) => stack = new_stack
             }
         }
@@ -693,16 +693,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_val, mut underflow) = unsafe {
+                let key_val = self.top.from_raw_mut().remove_as_leaf();
                 let underflow = self.top.from_raw().node().is_underfull();
-                (value, underflow)
+                (key_val, underflow)
             };
 
             loop {
@@ -717,7 +717,7 @@ mod stack {
                             self.map.depth -= 1;
                             self.map.root.hoist_lone_child();
                         }
-                        return value;
+                        return key_val;
                     }
                     Some(mut handle) => {
                         if underflow {
@@ -728,7 +728,7 @@ mod stack {
                             }
                         } else {
                             // All done!
-                            return value;
+                            return key_val;
                         }
                     }
                 }
@@ -739,7 +739,7 @@ mod stack {
     impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::LeafOrInternal> {
         /// 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 {
+        pub fn 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
@@ -1208,7 +1208,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
     /// Takes the value of the entry out of the map, and returns it.
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn remove(self) -> V {
-        self.stack.remove()
+        self.stack.remove().1
     }
 }
 
@@ -1609,3 +1609,86 @@ impl<K: Ord, V> BTreeMap<K, V> {
         }
     }
 }
+
+impl<K, Q: ?Sized> super::Recover<Q> for BTreeMap<K, ()> where K: Borrow<Q> + Ord, Q: Ord {
+    type Key = K;
+
+    fn get(&self, key: &Q) -> Option<&K> {
+        let mut cur_node = &self.root;
+        loop {
+            match Node::search(cur_node, key) {
+                Found(handle) => return Some(handle.into_kv().0),
+                GoDown(handle) => match handle.force() {
+                    Leaf(_) => return None,
+                    Internal(internal_handle) => {
+                        cur_node = internal_handle.into_edge();
+                        continue;
+                    }
+                }
+            }
+        }
+    }
+
+    fn take(&mut self, key: &Q) -> Option<K> {
+        // See `remove` for an explanation of this.
+
+        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()))
+                    },
+                    GoDown(handle) => {
+                        // We need to keep searching, try to go down the next edge
+                        match handle.force() {
+                            // We're at a leaf; the key isn't in here
+                            Leaf(_) => Finished(None),
+                            Internal(internal_handle) => Continue(pusher.push(internal_handle))
+                        }
+                    }
+                }
+            });
+            match result {
+                Finished(ret) => return ret.map(|(k, _)| k),
+                Continue(new_stack) => stack = new_stack
+            }
+        }
+    }
+
+    fn replace(&mut self, mut key: K) -> Option<K> {
+        // See `insert` for an explanation of this.
+
+        let mut stack = stack::PartialSearchStack::new(self);
+
+        loop {
+            let result = stack.with(move |pusher, node| {
+                match Node::search::<K, _>(node, &key) {
+                    Found(mut handle) => {
+                        mem::swap(handle.key_mut(), &mut key);
+                        Finished(Some(key))
+                    },
+                    GoDown(handle) => {
+                        match handle.force() {
+                            Leaf(leaf_handle) => {
+                                pusher.seal(leaf_handle).insert(key, ());
+                                Finished(None)
+                            }
+                            Internal(internal_handle) => {
+                                Continue((pusher.push(internal_handle), key, ()))
+                            }
+                        }
+                    }
+                }
+            });
+            match result {
+                Finished(ret) => return ret,
+                Continue((new_stack, renewed_key, _)) => {
+                    stack = new_stack;
+                    key = renewed_key;
+                }
+            }
+        }
+    }
+}

src/libcollections/btree/mod.rs

@@ -11,3 +11,11 @@
 mod node;
 pub mod map;
 pub mod set;
+
+trait Recover<Q: ?Sized> {
+    type Key;
+
+    fn get(&self, key: &Q) -> Option<&Self::Key>;
+    fn take(&mut self, key: &Q) -> Option<Self::Key>;
+    fn replace(&mut self, key: Self::Key) -> Option<Self::Key>;
+}

src/libcollections/btree/set.rs

@@ -19,6 +19,7 @@ use core::ops::{BitOr, BitAnd, BitXor, Sub};
 
 use borrow::Borrow;
 use btree_map::{BTreeMap, Keys};
+use super::Recover;
 use Bound;
 
 // FIXME(conventions): implement bounded iterators
@@ -329,6 +330,16 @@ impl<T: Ord> BTreeSet<T> {
         self.map.contains_key(value)
     }
 
+    /// Returns a reference to the value in the set, if any, that is equal to the given value.
+    ///
+    /// The value may be any borrowed form of the set's value type,
+    /// but the ordering on the borrowed form *must* match the
+    /// ordering on the value type.
+    #[unstable(feature = "set_recovery", issue = "28050")]
+    pub fn get<Q: ?Sized>(&self, value: &Q) -> Option<&T> where T: Borrow<Q>, Q: Ord {
+        Recover::get(&self.map, value)
+    }
+
     /// Returns `true` if the set has no elements in common with `other`.
     /// This is equivalent to checking for an empty intersection.
     ///
@@ -436,6 +447,13 @@ impl<T: Ord> BTreeSet<T> {
         self.map.insert(value, ()).is_none()
     }
 
+    /// Adds a value to the set, replacing the existing value, if any, that is equal to the given
+    /// one. Returns the replaced value.
+    #[unstable(feature = "set_recovery", issue = "28050")]
+    pub fn replace(&mut self, value: T) -> Option<T> {
+        Recover::replace(&mut self.map, value)
+    }
+
     /// Removes a value from the set. Returns `true` if the value was
     /// present in the set.
     ///
@@ -458,6 +476,16 @@ impl<T: Ord> BTreeSet<T> {
     pub fn remove<Q: ?Sized>(&mut self, value: &Q) -> bool where T: Borrow<Q>, Q: Ord {
         self.map.remove(value).is_some()
     }
+
+    /// Removes and returns the value in the set, if any, that is equal to the given one.
+    ///
+    /// The value may be any borrowed form of the set's value type,
+    /// but the ordering on the borrowed form *must* match the
+    /// ordering on the value type.
+    #[unstable(feature = "set_recovery", issue = "28050")]
+    pub fn take<Q: ?Sized>(&mut self, value: &Q) -> Option<T> where T: Borrow<Q>, Q: Ord {
+        Recover::take(&mut self.map, value)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

src/libcollectionstest/btree/set.rs

@@ -211,3 +211,52 @@ fn test_extend_ref() {
     assert!(a.contains(&5));
     assert!(a.contains(&6));
 }
+
+#[test]
+fn test_recovery() {
+    use std::cmp::Ordering;
+
+    #[derive(Debug)]
+    struct Foo(&'static str, i32);
+
+    impl PartialEq for Foo {
+        fn eq(&self, other: &Self) -> bool {
+            self.0 == other.0
+        }
+    }
+
+    impl Eq for Foo {}
+
+    impl PartialOrd for Foo {
+        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+            self.0.partial_cmp(&other.0)
+        }
+    }
+
+    impl Ord for Foo {
+        fn cmp(&self, other: &Self) -> Ordering {
+            self.0.cmp(&other.0)
+        }
+    }
+
+    let mut s = BTreeSet::new();
+    assert_eq!(s.replace(Foo("a", 1)), None);
+    assert_eq!(s.len(), 1);
+    assert_eq!(s.replace(Foo("a", 2)), Some(Foo("a", 1)));
+    assert_eq!(s.len(), 1);
+
+    {
+        let mut it = s.iter();
+        assert_eq!(it.next(), Some(&Foo("a", 2)));
+        assert_eq!(it.next(), None);
+    }
+
+    assert_eq!(s.get(&Foo("a", 1)), Some(&Foo("a", 2)));
+    assert_eq!(s.take(&Foo("a", 1)), Some(Foo("a", 2)));
+    assert_eq!(s.len(), 0);
+
+    assert_eq!(s.get(&Foo("a", 1)), None);
+    assert_eq!(s.take(&Foo("a", 1)), None);
+
+    assert_eq!(s.iter().next(), None);
+}

src/libcollectionstest/lib.rs

@@ -25,6 +25,7 @@
 #![feature(rand)]
 #![feature(range_inclusive)]
 #![feature(rustc_private)]
+#![feature(set_recovery)]
 #![feature(slice_bytes)]
 #![feature(slice_splits)]
 #![feature(split_off)]

src/libstd/collections/hash/map.rs

@@ -769,7 +769,7 @@ impl<K, V, S> HashMap<K, V, S>
     /// If the key already exists, the hashtable will be returned untouched
     /// and a reference to the existing element will be returned.
     fn insert_hashed_nocheck(&mut self, hash: SafeHash, k: K, v: V) -> &mut V {
-        self.insert_or_replace_with(hash, k, v, |_, _, _| ())
+        self.insert_or_replace_with(hash, k, v, |_, _, _, _| ())
     }
 
     fn insert_or_replace_with<'a, F>(&'a mut self,
@@ -778,7 +778,7 @@ impl<K, V, S> HashMap<K, V, S>
                                      v: V,
                                      mut found_existing: F)
                                      -> &'a mut V where
-        F: FnMut(&mut K, &mut V, V),
+        F: FnMut(&mut K, &mut V, K, V),
     {
         // Worst case, we'll find one empty bucket among `size + 1` buckets.
         let size = self.table.size();
@@ -801,7 +801,7 @@ impl<K, V, S> HashMap<K, V, S>
                     let (bucket_k, bucket_v) = bucket.into_mut_refs();
                     debug_assert!(k == *bucket_k);
                     // Key already exists. Get its reference.
-                    found_existing(bucket_k, bucket_v, v);
+                    found_existing(bucket_k, bucket_v, k, v);
                     return bucket_v;
                 }
             }
@@ -1123,7 +1123,7 @@ impl<K, V, S> HashMap<K, V, S>
         self.reserve(1);
 
         let mut retval = None;
-        self.insert_or_replace_with(hash, k, v, |_, val_ref, val| {
+        self.insert_or_replace_with(hash, k, v, |_, val_ref, _, val| {
             retval = Some(replace(val_ref, val));
         });
         retval
@@ -1630,6 +1630,35 @@ impl Default for RandomState {
     }
 }
 
+impl<K, S, Q: ?Sized> super::Recover<Q> for HashMap<K, (), S>
+    where K: Eq + Hash + Borrow<Q>, S: HashState, Q: Eq + Hash
+{
+    type Key = K;
+
+    fn get(&self, key: &Q) -> Option<&K> {
+        self.search(key).map(|bucket| bucket.into_refs().0)
+    }
+
+    fn take(&mut self, key: &Q) -> Option<K> {
+        if self.table.size() == 0 {
+            return None
+        }
+
+        self.search_mut(key).map(|bucket| pop_internal(bucket).0)
+    }
+
+    fn replace(&mut self, key: K) -> Option<K> {
+        let hash = self.make_hash(&key);
+        self.reserve(1);
+
+        let mut retkey = None;
+        self.insert_or_replace_with(hash, key, (), |key_ref, _, key, _| {
+            retkey = Some(replace(key_ref, key));
+        });
+        retkey
+    }
+}
+
 #[cfg(test)]
 mod test_map {
     use prelude::v1::*;

src/libstd/collections/hash/mod.rs

@@ -15,3 +15,11 @@ mod table;
 pub mod map;
 pub mod set;
 pub mod state;
+
+trait Recover<Q: ?Sized> {
+    type Key;
+
+    fn get(&self, key: &Q) -> Option<&Self::Key>;
+    fn take(&mut self, key: &Q) -> Option<Self::Key>;
+    fn replace(&mut self, key: Self::Key) -> Option<Self::Key>;
+}