libcollections: generalize BTree* to use BorrowFrom

Generalizes the BTree-based collections to use the new BorrowFrom
infrastructure for more flexible lookups and removals.

3 files changed, 42 insertions, 15 deletions

diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs
index 0b6d9e356b9..007e01ef982 100644
--- a/src/libcollections/btree/map.rs
+++ b/src/libcollections/btree/map.rs
@@ -21,6 +21,7 @@ use core::prelude::*;
 
 use self::StackOp::*;
 use super::node::*;
+use core::borrow::BorrowFrom;
 use std::hash::{Writer, Hash};
 use core::default::Default;
 use core::{iter, fmt, mem};
@@ -184,6 +185,9 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
     /// Returns a 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.
+    ///
     /// # Example
     ///
     /// ```
@@ -195,7 +199,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// assert_eq!(map.get(&2), None);
     /// ```
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
-    pub fn get(&self, key: &K) -> Option<&V> {
+    pub fn get<Sized? Q>(&self, key: &Q) -> Option<&V> where Q: BorrowFrom<K> + Ord {
         let mut cur_node = &self.root;
         loop {
             match cur_node.search(key) {
@@ -213,6 +217,9 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
     /// Returns true if the map contains a value for the specified 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.
+    ///
     /// # Example
     ///
     /// ```
@@ -224,7 +231,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// assert_eq!(map.contains_key(&2), false);
     /// ```
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
-    pub fn contains_key(&self, key: &K) -> bool {
+    pub fn contains_key<Sized? Q>(&self, key: &Q) -> bool where Q: BorrowFrom<K> + Ord {
         self.get(key).is_some()
     }
 
@@ -236,6 +243,9 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
     /// 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.
+    ///
     /// # Example
     ///
     /// ```
@@ -251,7 +261,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// ```
     // See `get` for implementation notes, this is basically a copy-paste with mut's added
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
-    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+    pub fn get_mut<Sized? Q>(&mut self, key: &Q) -> Option<&mut V> where Q: BorrowFrom<K> + Ord {
         // temp_node is a Borrowck hack for having a mutable value outlive a loop iteration
         let mut temp_node = &mut self.root;
         loop {
@@ -410,6 +420,9 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// Removes a key from the map, returning 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.
+    ///
     /// # Example
     ///
     /// ```
@@ -421,7 +434,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// assert_eq!(map.remove(&1), None);
     /// ```
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
-    pub fn remove(&mut self, key: &K) -> Option<V> {
+    pub fn remove<Sized? Q>(&mut self, key: &Q) -> Option<V> where Q: BorrowFrom<K> + Ord {
         // See `swap` for a more thorough description of the stuff going on in here
         let mut stack = stack::PartialSearchStack::new(self);
         loop {
@@ -790,14 +803,18 @@ impl<K: Show, V: Show> Show for BTreeMap<K, V> {
     }
 }
 
-impl<K: Ord, V> Index<K, V> for BTreeMap<K, V> {
-    fn index(&self, key: &K) -> &V {
+impl<K: Ord, Sized? Q, V> Index<Q, V> for BTreeMap<K, V>
+    where Q: BorrowFrom<K> + Ord
+{
+    fn index(&self, key: &Q) -> &V {
         self.get(key).expect("no entry found for key")
     }
 }
 
-impl<K: Ord, V> IndexMut<K, V> for BTreeMap<K, V> {
-    fn index_mut(&mut self, key: &K) -> &mut V {
+impl<K: Ord, Sized? Q, V> IndexMut<Q, V> for BTreeMap<K, V>
+    where Q: BorrowFrom<K> + Ord
+{
+    fn index_mut(&mut self, key: &Q) -> &mut V {
         self.get_mut(key).expect("no entry found for key")
     }
 }
diff --git a/src/libcollections/btree/node.rs b/src/libcollections/btree/node.rs
index 6e9341231ad..bdd7aa9c611 100644
--- a/src/libcollections/btree/node.rs
+++ b/src/libcollections/btree/node.rs
@@ -19,6 +19,7 @@ use core::prelude::*;
 
 use core::{slice, mem, ptr};
 use core::iter::Zip;
+use core::borrow::BorrowFrom;
 
 use vec;
 use vec::Vec;
@@ -73,19 +74,19 @@ impl<K: Ord, V> Node<K, V> {
     /// Searches for the given key in the node. If it finds an exact match,
     /// `Found` will be yielded with the matching index. If it doesn't find an exact match,
     /// `GoDown` will be yielded with the index of the subtree the key must lie in.
-    pub fn search(&self, key: &K) -> SearchResult {
+    pub fn search<Sized? Q>(&self, key: &Q) -> SearchResult where Q: BorrowFrom<K> + Ord {
         // FIXME(Gankro): Tune when to search linear or binary based on B (and maybe K/V).
         // For the B configured as of this writing (B = 6), binary search was *significantly*
         // worse for uints.
         self.search_linear(key)
     }
 
-    fn search_linear(&self, key: &K) -> SearchResult {
+    fn search_linear<Sized? Q>(&self, key: &Q) -> SearchResult where Q: BorrowFrom<K> + Ord {
         for (i, k) in self.keys.iter().enumerate() {
-            match k.cmp(key) {
-                Less => {},
+            match key.cmp(BorrowFrom::borrow_from(k)) {
+                Greater => {},
                 Equal => return Found(i),
-                Greater => return GoDown(i),
+                Less => return GoDown(i),
             }
         }
         GoDown(self.len())
diff --git a/src/libcollections/btree/set.rs b/src/libcollections/btree/set.rs
index 56edf7719bb..64ae4f6a508 100644
--- a/src/libcollections/btree/set.rs
+++ b/src/libcollections/btree/set.rs
@@ -15,6 +15,7 @@ use core::prelude::*;
 
 use btree_map::{BTreeMap, Keys, MoveEntries};
 use std::hash::Hash;
+use core::borrow::BorrowFrom;
 use core::default::Default;
 use core::{iter, fmt};
 use core::iter::Peekable;
@@ -167,6 +168,10 @@ impl<T: Ord> BTreeSet<T> {
 
     /// Returns `true` if the set contains a 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.
+    ///
     /// # Example
     ///
     /// ```
@@ -177,7 +182,7 @@ impl<T: Ord> BTreeSet<T> {
     /// assert_eq!(set.contains(&4), false);
     /// ```
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
-    pub fn contains(&self, value: &T) -> bool {
+    pub fn contains<Sized? Q>(&self, value: &Q) -> bool where Q: BorrowFrom<T> + Ord {
         self.map.contains_key(value)
     }
 
@@ -291,6 +296,10 @@ impl<T: Ord> BTreeSet<T> {
     /// Removes a value from the set. Returns `true` if the value was
     /// present in the set.
     ///
+    /// 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.
+    ///
     /// # Example
     ///
     /// ```
@@ -303,7 +312,7 @@ impl<T: Ord> BTreeSet<T> {
     /// assert_eq!(set.remove(&2), false);
     /// ```
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
-    pub fn remove(&mut self, value: &T) -> bool {
+    pub fn remove<Sized? Q>(&mut self, value: &Q) -> bool where Q: BorrowFrom<T> + Ord {
         self.map.remove(value).is_some()
     }
 }