diff --git a/src/libcollections/treemap.rs b/src/libcollections/treemap.rs index 4a7fcd6786fa0..e0242af3c9973 100644 --- a/src/libcollections/treemap.rs +++ b/src/libcollections/treemap.rs @@ -88,40 +88,18 @@ impl Mutable for TreeMap { } impl Map for TreeMap { + // See comments on tree_find_with + #[inline] fn find<'a>(&'a self, key: &K) -> Option<&'a V> { - let mut current = &self.root; - loop { - match *current { - Some(ref r) => { - match key.cmp(&r.key) { - Less => current = &r.left, - Greater => current = &r.right, - Equal => return Some(&r.value) - } - } - None => return None - } - } + tree_find_with(&self.root, |k2| key.cmp(k2)) } } impl MutableMap for TreeMap { + // See comments on def_tree_find_mut_with #[inline] fn find_mut<'a>(&'a mut self, key: &K) -> Option<&'a mut V> { - let mut current = &mut self.root; - loop { - let temp = current; // hack to appease borrowck - match *temp { - Some(ref mut r) => { - match key.cmp(&r.key) { - Less => current = &mut r.left, - Greater => current = &mut r.right, - Equal => return Some(&mut r.value) - } - } - None => return None - } - } + tree_find_mut_with(&mut self.root, |x| key.cmp(x)) } fn swap(&mut self, key: K, value: V) -> Option { @@ -194,6 +172,55 @@ impl TreeMap { } } +impl TreeMap { + /// Return the value for which f(key) returns Equal. f is invoked + /// with current key and helps to navigate the tree + /// + /// # Example + /// + /// ``` + /// use std::ascii::StrAsciiExt; + /// + /// let mut t = collections::treemap::TreeMap::new(); + /// t.insert("Content-Type", "application/xml"); + /// t.insert("User-Agent", "Curl-Rust/0.1"); + /// + /// let ua_key = "user-agent"; + /// let ua = t.find_with(|&k| { + /// ua_key.cmp(&k.to_ascii_lower().as_slice()) + /// }); + /// + /// assert_eq!(*ua.unwrap(), "Curl-Rust/0.1"); + /// ``` + #[inline] + pub fn find_with<'a>(&'a self, f:|&K| -> Ordering) -> Option<&'a V> { + tree_find_with(&self.root, f) + } + + /// Return the value for which f(key) returns Equal. f is invoked + /// with current key and helps to navigate the tree + /// + /// # Example + /// + /// ``` + /// let mut t = collections::treemap::TreeMap::new(); + /// t.insert("Content-Type", "application/xml"); + /// t.insert("User-Agent", "Curl-Rust/0.1"); + /// + /// let new_ua = "Safari/156.0"; + /// match t.find_mut_with(|k| "User-Agent".cmp(k)) { + /// Some(x) => *x = new_ua, + /// None => fail!(), + /// } + /// + /// assert_eq!(t.find(&"User-Agent"), Some(&new_ua)); + /// ``` + #[inline] + pub fn find_mut_with<'a>(&'a mut self, f:|&K| -> Ordering) -> Option<&'a mut V> { + tree_find_mut_with(&mut self.root, f) + } +} + // range iterators. macro_rules! bound_setup { @@ -853,6 +880,51 @@ fn split(node: &mut Box>) { } } +// Next 2 functions have the same conventions +// +// The only difference is that non-mutable version uses loop instead +// of recursion (performance considerations) +// It seems to be impossible to avoid recursion with mutability +// +// So convention is that comparator is gets at input current key +// and returns search_key cmp cur_key (i.e. search_key.cmp(cur_key)) +fn tree_find_with<'r, K, V>(node: &'r Option>>, + f: |&K| -> Ordering) -> Option<&'r V> { + let mut current: &'r Option>> = node; + loop { + match *current { + Some(ref r) => { + match f(&r.key) { + Less => current = &r.left, + Greater => current = &r.right, + Equal => return Some(&r.value) + } + } + None => return None + } + } +} + +// See comments above tree_find_with +fn tree_find_mut_with<'r, K, V>(node: &'r mut Option>>, + f: |&K| -> Ordering) -> Option<&'r mut V> { + + let mut current = node; + loop { + let temp = current; // hack to appease borrowck + match *temp { + Some(ref mut r) => { + match f(&r.key) { + Less => current = &mut r.left, + Greater => current = &mut r.right, + Equal => return Some(&mut r.value) + } + } + None => return None + } + } +} + fn insert(node: &mut Option>>, key: K, value: V) -> Option { match *node { @@ -1024,6 +1096,30 @@ mod test_treemap { assert_eq!(m.find(&2), None); } + #[test] + fn find_with_empty() { + let m: TreeMap<&'static str,int> = TreeMap::new(); + assert!(m.find_with(|k| "test".cmp(k)) == None); + } + + #[test] + fn find_with_not_found() { + let mut m = TreeMap::new(); + assert!(m.insert("test1", 2i)); + assert!(m.insert("test2", 3i)); + assert!(m.insert("test3", 3i)); + assert_eq!(m.find_with(|k| "test4".cmp(k)), None); + } + + #[test] + fn find_with_found() { + let mut m = TreeMap::new(); + assert!(m.insert("test1", 2i)); + assert!(m.insert("test2", 3i)); + assert!(m.insert("test3", 4i)); + assert_eq!(m.find_with(|k| "test2".cmp(k)), Some(&3i)); + } + #[test] fn test_find_mut() { let mut m = TreeMap::new(); @@ -1037,6 +1133,19 @@ mod test_treemap { assert_eq!(m.find(&5), Some(&new)); } + #[test] + fn test_find_with_mut() { + let mut m = TreeMap::new(); + assert!(m.insert("t1", 12i)); + assert!(m.insert("t2", 8)); + assert!(m.insert("t5", 14)); + let new = 100; + match m.find_mut_with(|k| "t5".cmp(k)) { + None => fail!(), Some(x) => *x = new + } + assert_eq!(m.find_with(|k| "t5".cmp(k)), Some(&new)); + } + #[test] fn insert_replace() { let mut m = TreeMap::new();