Skip to content

Commit f9d3b9e

Browse files
committed
auto merge of #15220 : vhbit/rust/treemap-str-equiv, r=alexcrichton
- it allows to lookup using any str-equiv object, making TreeMaps finally usable (for example, it is much easier to work with JSON with lookup values being static strs) - actually provides pretty flexible solution which could be extended to other equivalent types (although it might be not that performant)
2 parents 3f3291e + be7a170 commit f9d3b9e

File tree

1 file changed

+136
-27
lines changed

1 file changed

+136
-27
lines changed

src/libcollections/treemap.rs

+136-27
Original file line numberDiff line numberDiff line change
@@ -88,40 +88,18 @@ impl<K: Ord, V> Mutable for TreeMap<K, V> {
8888
}
8989

9090
impl<K: Ord, V> Map<K, V> for TreeMap<K, V> {
91+
// See comments on tree_find_with
92+
#[inline]
9193
fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
92-
let mut current = &self.root;
93-
loop {
94-
match *current {
95-
Some(ref r) => {
96-
match key.cmp(&r.key) {
97-
Less => current = &r.left,
98-
Greater => current = &r.right,
99-
Equal => return Some(&r.value)
100-
}
101-
}
102-
None => return None
103-
}
104-
}
94+
tree_find_with(&self.root, |k2| key.cmp(k2))
10595
}
10696
}
10797

10898
impl<K: Ord, V> MutableMap<K, V> for TreeMap<K, V> {
99+
// See comments on def_tree_find_mut_with
109100
#[inline]
110101
fn find_mut<'a>(&'a mut self, key: &K) -> Option<&'a mut V> {
111-
let mut current = &mut self.root;
112-
loop {
113-
let temp = current; // hack to appease borrowck
114-
match *temp {
115-
Some(ref mut r) => {
116-
match key.cmp(&r.key) {
117-
Less => current = &mut r.left,
118-
Greater => current = &mut r.right,
119-
Equal => return Some(&mut r.value)
120-
}
121-
}
122-
None => return None
123-
}
124-
}
102+
tree_find_mut_with(&mut self.root, |x| key.cmp(x))
125103
}
126104

127105
fn swap(&mut self, key: K, value: V) -> Option<V> {
@@ -194,6 +172,55 @@ impl<K: Ord, V> TreeMap<K, V> {
194172
}
195173
}
196174

175+
impl<K, V> TreeMap<K, V> {
176+
/// Return the value for which f(key) returns Equal. f is invoked
177+
/// with current key and helps to navigate the tree
178+
///
179+
/// # Example
180+
///
181+
/// ```
182+
/// use std::ascii::StrAsciiExt;
183+
///
184+
/// let mut t = collections::treemap::TreeMap::new();
185+
/// t.insert("Content-Type", "application/xml");
186+
/// t.insert("User-Agent", "Curl-Rust/0.1");
187+
///
188+
/// let ua_key = "user-agent";
189+
/// let ua = t.find_with(|&k| {
190+
/// ua_key.cmp(&k.to_ascii_lower().as_slice())
191+
/// });
192+
///
193+
/// assert_eq!(*ua.unwrap(), "Curl-Rust/0.1");
194+
/// ```
195+
#[inline]
196+
pub fn find_with<'a>(&'a self, f:|&K| -> Ordering) -> Option<&'a V> {
197+
tree_find_with(&self.root, f)
198+
}
199+
200+
/// Return the value for which f(key) returns Equal. f is invoked
201+
/// with current key and helps to navigate the tree
202+
///
203+
/// # Example
204+
///
205+
/// ```
206+
/// let mut t = collections::treemap::TreeMap::new();
207+
/// t.insert("Content-Type", "application/xml");
208+
/// t.insert("User-Agent", "Curl-Rust/0.1");
209+
///
210+
/// let new_ua = "Safari/156.0";
211+
/// match t.find_mut_with(|k| "User-Agent".cmp(k)) {
212+
/// Some(x) => *x = new_ua,
213+
/// None => fail!(),
214+
/// }
215+
///
216+
/// assert_eq!(t.find(&"User-Agent"), Some(&new_ua));
217+
/// ```
218+
#[inline]
219+
pub fn find_mut_with<'a>(&'a mut self, f:|&K| -> Ordering) -> Option<&'a mut V> {
220+
tree_find_mut_with(&mut self.root, f)
221+
}
222+
}
223+
197224
// range iterators.
198225

199226
macro_rules! bound_setup {
@@ -853,6 +880,51 @@ fn split<K: Ord, V>(node: &mut Box<TreeNode<K, V>>) {
853880
}
854881
}
855882

883+
// Next 2 functions have the same conventions
884+
//
885+
// The only difference is that non-mutable version uses loop instead
886+
// of recursion (performance considerations)
887+
// It seems to be impossible to avoid recursion with mutability
888+
//
889+
// So convention is that comparator is gets at input current key
890+
// and returns search_key cmp cur_key (i.e. search_key.cmp(cur_key))
891+
fn tree_find_with<'r, K, V>(node: &'r Option<Box<TreeNode<K, V>>>,
892+
f: |&K| -> Ordering) -> Option<&'r V> {
893+
let mut current: &'r Option<Box<TreeNode<K, V>>> = node;
894+
loop {
895+
match *current {
896+
Some(ref r) => {
897+
match f(&r.key) {
898+
Less => current = &r.left,
899+
Greater => current = &r.right,
900+
Equal => return Some(&r.value)
901+
}
902+
}
903+
None => return None
904+
}
905+
}
906+
}
907+
908+
// See comments above tree_find_with
909+
fn tree_find_mut_with<'r, K, V>(node: &'r mut Option<Box<TreeNode<K, V>>>,
910+
f: |&K| -> Ordering) -> Option<&'r mut V> {
911+
912+
let mut current = node;
913+
loop {
914+
let temp = current; // hack to appease borrowck
915+
match *temp {
916+
Some(ref mut r) => {
917+
match f(&r.key) {
918+
Less => current = &mut r.left,
919+
Greater => current = &mut r.right,
920+
Equal => return Some(&mut r.value)
921+
}
922+
}
923+
None => return None
924+
}
925+
}
926+
}
927+
856928
fn insert<K: Ord, V>(node: &mut Option<Box<TreeNode<K, V>>>,
857929
key: K, value: V) -> Option<V> {
858930
match *node {
@@ -1024,6 +1096,30 @@ mod test_treemap {
10241096
assert_eq!(m.find(&2), None);
10251097
}
10261098

1099+
#[test]
1100+
fn find_with_empty() {
1101+
let m: TreeMap<&'static str,int> = TreeMap::new();
1102+
assert!(m.find_with(|k| "test".cmp(k)) == None);
1103+
}
1104+
1105+
#[test]
1106+
fn find_with_not_found() {
1107+
let mut m = TreeMap::new();
1108+
assert!(m.insert("test1", 2i));
1109+
assert!(m.insert("test2", 3i));
1110+
assert!(m.insert("test3", 3i));
1111+
assert_eq!(m.find_with(|k| "test4".cmp(k)), None);
1112+
}
1113+
1114+
#[test]
1115+
fn find_with_found() {
1116+
let mut m = TreeMap::new();
1117+
assert!(m.insert("test1", 2i));
1118+
assert!(m.insert("test2", 3i));
1119+
assert!(m.insert("test3", 4i));
1120+
assert_eq!(m.find_with(|k| "test2".cmp(k)), Some(&3i));
1121+
}
1122+
10271123
#[test]
10281124
fn test_find_mut() {
10291125
let mut m = TreeMap::new();
@@ -1037,6 +1133,19 @@ mod test_treemap {
10371133
assert_eq!(m.find(&5), Some(&new));
10381134
}
10391135

1136+
#[test]
1137+
fn test_find_with_mut() {
1138+
let mut m = TreeMap::new();
1139+
assert!(m.insert("t1", 12i));
1140+
assert!(m.insert("t2", 8));
1141+
assert!(m.insert("t5", 14));
1142+
let new = 100;
1143+
match m.find_mut_with(|k| "t5".cmp(k)) {
1144+
None => fail!(), Some(x) => *x = new
1145+
}
1146+
assert_eq!(m.find_with(|k| "t5".cmp(k)), Some(&new));
1147+
}
1148+
10401149
#[test]
10411150
fn insert_replace() {
10421151
let mut m = TreeMap::new();

0 commit comments

Comments
 (0)