diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs index 1cf980b105941..f4fd81437fc6e 100644 --- a/src/libextra/treemap.rs +++ b/src/libextra/treemap.rs @@ -12,10 +12,10 @@ //! trees. The only requirement for the types is that the key implements //! `TotalOrd`. - use std::util::{swap, replace}; use std::iter::{Peekable}; use std::cmp::Ordering; +use std::ptr; // This is implemented as an AA tree, which is a simplified variation of // a red-black tree where red (horizontal) nodes can only be added @@ -135,11 +135,6 @@ impl TreeMap { /// Create an empty TreeMap pub fn new() -> TreeMap { TreeMap{root: None, length: 0} } - /// Iterate over the map and mutate the contained values - pub fn mutate_values(&mut self, f: |&K, &mut V| -> bool) -> bool { - mutate_values(&mut self.root, f) - } - /// Get a lazy iterator over the key-value pairs in the map. /// Requires that it be frozen (immutable). pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> { @@ -157,8 +152,74 @@ impl TreeMap { TreeMapRevIterator{iter: self.iter()} } + /// Get a lazy forward iterator over the key-value pairs in the + /// map, with the values being mutable. + pub fn mut_iter<'a>(&'a mut self) -> TreeMapMutIterator<'a, K, V> { + TreeMapMutIterator { + stack: ~[], + node: mut_deref(&mut self.root), + remaining_min: self.length, + remaining_max: self.length + } + } + /// Get a lazy reverse iterator over the key-value pairs in the + /// map, with the values being mutable. + pub fn mut_rev_iter<'a>(&'a mut self) -> TreeMapMutRevIterator<'a, K, V> { + TreeMapMutRevIterator{iter: self.mut_iter()} + } + + + /// Get a lazy iterator that consumes the treemap. + pub fn move_iter(self) -> TreeMapMoveIterator { + let TreeMap { root: root, length: length } = self; + let stk = match root { + None => ~[], + Some(~tn) => ~[tn] + }; + TreeMapMoveIterator { + stack: stk, + remaining: length + } + } +} + +// range iterators. + +macro_rules! bound_setup { + // initialiser of the iterator to manipulate + ($iter:expr, + // whether we are looking for the lower or upper bound. + $is_lower_bound:expr) => { + { + let mut iter = $iter; + loop { + if !iter.node.is_null() { + let node_k = unsafe {&(*iter.node).key}; + match k.cmp(node_k) { + Less => iter.traverse_left(), + Greater => iter.traverse_right(), + Equal => { + if $is_lower_bound { + iter.traverse_complete(); + return iter; + } else { + iter.traverse_right() + } + } + } + } else { + iter.traverse_complete(); + return iter; + } + } + } + } +} + + +impl TreeMap { /// Get a lazy iterator that should be initialized using - /// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`. + /// `traverse_left`/`traverse_right`/`traverse_complete`. fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> { TreeMapIterator { stack: ~[], @@ -171,175 +232,257 @@ impl TreeMap { /// Return a lazy iterator to the first key-value pair whose key is not less than `k` /// If all keys in map are less than `k` an empty iterator is returned. pub fn lower_bound<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> { - let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal(); - loop { - match iter.node { - Some(r) => { - match k.cmp(&r.key) { - Less => iter_traverse_left(&mut iter), - Greater => iter_traverse_right(&mut iter), - Equal => { - iter_traverse_complete(&mut iter); - return iter; - } - } - } - None => { - iter_traverse_complete(&mut iter); - return iter; - } - } - } + bound_setup!(self.iter_for_traversal(), true) } /// Return a lazy iterator to the first key-value pair whose key is greater than `k` /// If all keys in map are not greater than `k` an empty iterator is returned. pub fn upper_bound<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> { - let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal(); - loop { - match iter.node { - Some(r) => { - match k.cmp(&r.key) { - Less => iter_traverse_left(&mut iter), - Greater => iter_traverse_right(&mut iter), - Equal => iter_traverse_right(&mut iter) - } - } - None => { - iter_traverse_complete(&mut iter); - return iter; - } - } - } + bound_setup!(self.iter_for_traversal(), false) } - /// Get a lazy iterator that consumes the treemap. - pub fn move_iter(self) -> TreeMapMoveIterator { - let TreeMap { root: root, length: length } = self; - let stk = match root { - None => ~[], - Some(~tn) => ~[tn] - }; - TreeMapMoveIterator { - stack: stk, - remaining: length + /// Get a lazy iterator that should be initialized using + /// `traverse_left`/`traverse_right`/`traverse_complete`. + fn mut_iter_for_traversal<'a>(&'a mut self) -> TreeMapMutIterator<'a, K, V> { + TreeMapMutIterator { + stack: ~[], + node: mut_deref(&mut self.root), + remaining_min: 0, + remaining_max: self.length } } + + /// Return a lazy value iterator to the first key-value pair (with + /// the value being mutable) whose key is not less than `k`. + /// + /// If all keys in map are less than `k` an empty iterator is + /// returned. + pub fn mut_lower_bound<'a>(&'a mut self, k: &K) -> TreeMapMutIterator<'a, K, V> { + bound_setup!(self.mut_iter_for_traversal(), true) + } + + /// Return a lazy iterator to the first key-value pair (with the + /// value being mutable) whose key is greater than `k`. + /// + /// If all keys in map are not greater than `k` an empty iterator + /// is returned. + pub fn mut_upper_bound<'a>(&'a mut self, k: &K) -> TreeMapMutIterator<'a, K, V> { + bound_setup!(self.mut_iter_for_traversal(), false) + } } /// Lazy forward iterator over a map pub struct TreeMapIterator<'a, K, V> { priv stack: ~[&'a TreeNode], - priv node: Option<&'a TreeNode>, + // See the comment on TreeMapMutIterator; this is just to allow + // code-sharing (for this immutable-values iterator it *could* very + // well be Option<&'a TreeNode>). + priv node: *TreeNode, priv remaining_min: uint, priv remaining_max: uint } -fn deref<'a, K, V>(node: &'a Option<~TreeNode>) -> Option<&'a TreeNode> { - match *node { - Some(ref n) => { - let n: &TreeNode = *n; - Some(n) - } - None => None - } +/// Lazy backward iterator over a map +pub struct TreeMapRevIterator<'a, K, V> { + priv iter: TreeMapIterator<'a, K, V>, } -impl<'a, K, V> TreeMapIterator<'a, K, V> { - #[inline(always)] - fn next_(&mut self, forward: bool) -> Option<(&'a K, &'a V)> { - while !self.stack.is_empty() || self.node.is_some() { - match self.node { - Some(x) => { - self.stack.push(x); - self.node = deref(if forward { &x.left } else { &x.right }); - } - None => { - let res = self.stack.pop(); - self.node = deref(if forward { &res.right } else { &res.left }); - self.remaining_max -= 1; - if self.remaining_min > 0 { - self.remaining_min -= 1; +/// Lazy forward iterator over a map that allows for the mutation of +/// the values. +pub struct TreeMapMutIterator<'a, K, V> { + priv stack: ~[&'a mut TreeNode], + // Unfortunately, we require some unsafe-ness to get around the + // fact that we would be storing a reference *into* one of the + // nodes in the stack. + // + // As far as the compiler knows, this would let us invalidate the + // reference by assigning a new value to this node's position in + // its parent, which would cause this current one to be + // deallocated so this reference would be invalid. (i.e. the + // compilers complaints are 100% correct.) + // + // However, as far as you humans reading this code know (or are + // about to know, if you haven't read far enough down yet), we are + // only reading from the TreeNode.{left,right} fields. the only + // thing that is ever mutated is the .value field (although any + // actual mutation that happens is done externally, by the + // iterator consumer). So, don't be so concerned, rustc, we've got + // it under control. + // + // (This field can legitimately be null.) + priv node: *mut TreeNode, + priv remaining_min: uint, + priv remaining_max: uint +} + +/// Lazy backward iterator over a map +pub struct TreeMapMutRevIterator<'a, K, V> { + priv iter: TreeMapMutIterator<'a, K, V>, +} + + +// FIXME #5846 we want to be able to choose between &x and &mut x +// (with many different `x`) below, so we need to optionally pass mut +// as a tt, but the only thing we can do with a `tt` is pass them to +// other macros, so this takes the `& ` token +// sequence and forces their evalutation as an expression. +macro_rules! addr { ($e:expr) => { $e }} + +macro_rules! define_iterator { + ($name:ident, + $rev_name:ident, + // the type of the values of the treemap in the return value of + // the iterator (i.e. &V or &mut V). This is non-hygienic in the + // name of the lifetime. + value_type = $value_type:ty, + + // the function to go from &m Option<~TreeNode> to *m TreeNode + deref = $deref:ident, + + // see comment on `addr!`, this is just an optional `mut`, but + // there's no support for 0-or-1 repeats. + addr_mut = $($addr_mut:tt)* + ) => { + // private methods on the forward iterator + impl<'a, K, V> $name<'a, K, V> { + #[inline(always)] + fn next_(&mut self, forward: bool) -> Option<(&'a K, $value_type)> { + while !self.stack.is_empty() || !self.node.is_null() { + if !self.node.is_null() { + let node = unsafe {addr!(& $($addr_mut)* *self.node)}; + { + let next_node = if forward { + addr!(& $($addr_mut)* node.left) + } else { + addr!(& $($addr_mut)* node.right) + }; + self.node = $deref(next_node); + } + self.stack.push(node); + } else { + let node = self.stack.pop(); + let next_node = if forward { + addr!(& $($addr_mut)* node.right) + } else { + addr!(& $($addr_mut)* node.left) + }; + self.node = $deref(next_node); + self.remaining_max -= 1; + if self.remaining_min > 0 { + self.remaining_min -= 1; + } + return Some((&node.key, addr!(& $($addr_mut)* node.value))); + } + } + None + } + + /// traverse_left, traverse_right and traverse_complete are + /// used to initialize TreeMapIterator/TreeMapMutIterator + /// pointing to element inside tree structure. + /// + /// They should be used in following manner: + /// - create iterator using TreeMap::[mut_]iter_for_traversal + /// - find required node using `traverse_left`/`traverse_right` + /// (current node is `TreeMapIterator::node` field) + /// - complete initialization with `traverse_complete` + /// + /// After this, iteration will start from `self.node`. If + /// `self.node` is None iteration will start from last + /// node from which we traversed left. + #[inline] + fn traverse_left(&mut self) { + let node = unsafe {addr!(& $($addr_mut)* *self.node)}; + self.node = $deref(addr!(& $($addr_mut)* node.left)); + self.stack.push(node); + } + + #[inline] + fn traverse_right(&mut self) { + let node = unsafe {addr!(& $($addr_mut)* *self.node)}; + self.node = $deref(addr!(& $($addr_mut)* node.right)); + } + + #[inline] + fn traverse_complete(&mut self) { + if !self.node.is_null() { + unsafe { + self.stack.push(addr!(& $($addr_mut)* *self.node)); + } + self.node = ptr::RawPtr::null(); } - return Some((&res.key, &res.value)); - } } } - None - } -} -impl<'a, K, V> Iterator<(&'a K, &'a V)> for TreeMapIterator<'a, K, V> { - /// Advance the iterator to the next node (in order) and return a - /// tuple with a reference to the key and value. If there are no - /// more nodes, return `None`. - fn next(&mut self) -> Option<(&'a K, &'a V)> { - self.next_(true) - } + // the forward Iterator impl. + impl<'a, K, V> Iterator<(&'a K, $value_type)> for $name<'a, K, V> { + /// Advance the iterator to the next node (in order) and return a + /// tuple with a reference to the key and value. If there are no + /// more nodes, return `None`. + fn next(&mut self) -> Option<(&'a K, $value_type)> { + self.next_(true) + } - #[inline] - fn size_hint(&self) -> (uint, Option) { - (self.remaining_min, Some(self.remaining_max)) - } -} + #[inline] + fn size_hint(&self) -> (uint, Option) { + (self.remaining_min, Some(self.remaining_max)) + } + } -/// Lazy backward iterator over a map -pub struct TreeMapRevIterator<'a, K, V> { - priv iter: TreeMapIterator<'a, K, V>, -} + // the reverse Iterator impl. + impl<'a, K, V> Iterator<(&'a K, $value_type)> for $rev_name<'a, K, V> { + fn next(&mut self) -> Option<(&'a K, $value_type)> { + self.iter.next_(false) + } -impl<'a, K, V> Iterator<(&'a K, &'a V)> for TreeMapRevIterator<'a, K, V> { - /// Advance the iterator to the next node (in order) and return a - /// tuple with a reference to the key and value. If there are no - /// more nodes, return `None`. - fn next(&mut self) -> Option<(&'a K, &'a V)> { - self.iter.next_(false) + #[inline] + fn size_hint(&self) -> (uint, Option) { + self.iter.size_hint() + } + } } +} // end of define_iterator - #[inline] - fn size_hint(&self) -> (uint, Option) { - self.iter.size_hint() - } +define_iterator! { + TreeMapIterator, + TreeMapRevIterator, + value_type = &'a V, + deref = deref, + + // immutable, so no mut + addr_mut = } +define_iterator! { + TreeMapMutIterator, + TreeMapMutRevIterator, + value_type = &'a mut V, + deref = mut_deref, -/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to -/// initialize TreeMapIterator pointing to element inside tree structure. -/// -/// They should be used in following manner: -/// - create iterator using TreeMap::iter_for_traversal -/// - find required node using `iter_traverse_left`/`iter_traverse_right` -/// (current node is `TreeMapIterator::node` field) -/// - complete initialization with `iter_traverse_complete` -#[inline] -fn iter_traverse_left<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) { - let node = it.node.unwrap(); - it.stack.push(node); - it.node = deref(&node.left); + addr_mut = mut } -#[inline] -fn iter_traverse_right<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) { - it.node = deref(&it.node.get_ref().right); +fn deref<'a, K, V>(node: &'a Option<~TreeNode>) -> *TreeNode { + match *node { + Some(ref n) => { + let n: &TreeNode = *n; + n as *TreeNode + } + None => ptr::null() + } } -/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to -/// initialize TreeMapIterator pointing to element inside tree structure. -/// -/// Completes traversal. Should be called before using iterator. -/// Iteration will start from `self.node`. -/// If `self.node` is None iteration will start from last node from which we -/// traversed left. -#[inline] -fn iter_traverse_complete<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) { - match it.node { - Some(n) => { - it.stack.push(n); - it.node = None; +fn mut_deref(x: &mut Option<~TreeNode>) -> *mut TreeNode { + match *x { + Some(ref mut n) => { + let n: &mut TreeNode = *n; + n as *mut TreeNode } - None => () + None => ptr::mut_null() } } + + /// Lazy forward iterator over a map that consumes the map while iterating pub struct TreeMapMoveIterator { priv stack: ~[TreeNode], @@ -678,24 +821,6 @@ impl TreeNode { } } -fn mutate_values<'r, - K:TotalOrd, - V>( - node: &'r mut Option<~TreeNode>, - f: |&'r K, &'r mut V| -> bool) - -> bool { - match *node { - Some(~TreeNode{key: ref key, value: ref mut value, left: ref mut left, - right: ref mut right, ..}) => { - if !mutate_values(left, |k,v| f(k,v)) { return false } - if !f(key, value) { return false } - if !mutate_values(right, |k,v| f(k,v)) { return false } - } - None => return false - } - true -} - // Remove left horizontal link by rotating right fn skew(node: &mut ~TreeNode) { if node.left.as_ref().map_default(false, |x| x.level == node.level) { @@ -1129,6 +1254,69 @@ mod test_treemap { } } + #[test] + fn test_mut_iter() { + let mut m = TreeMap::new(); + for i in range(0u, 10) { + assert!(m.insert(i, 100 * i)); + } + + for (i, (&k, v)) in m.mut_iter().enumerate() { + *v += k * 10 + i; // 000 + 00 + 0, 100 + 10 + 1, ... + } + + for (&k, &v) in m.iter() { + assert_eq!(v, 111 * k); + } + } + #[test] + fn test_mut_rev_iter() { + let mut m = TreeMap::new(); + for i in range(0u, 10) { + assert!(m.insert(i, 100 * i)); + } + + for (i, (&k, v)) in m.mut_rev_iter().enumerate() { + *v += k * 10 + (9 - i); // 900 + 90 + (9 - 0), 800 + 80 + (9 - 1), ... + } + + for (&k, &v) in m.iter() { + assert_eq!(v, 111 * k); + } + } + + #[test] + fn test_mut_interval_iter() { + let mut m_lower = TreeMap::new(); + let mut m_upper = TreeMap::new(); + for i in range(1, 100) { + assert!(m_lower.insert(i * 2, i * 4)); + assert!(m_upper.insert(i * 2, i * 4)); + } + + for i in range(1, 199) { + let mut lb_it = m_lower.mut_lower_bound(&i); + let (&k, v) = lb_it.next().unwrap(); + let lb = i + i % 2; + assert_eq!(lb, k); + *v -= k; + } + for i in range(0, 198) { + let mut ub_it = m_upper.mut_upper_bound(&i); + let (&k, v) = ub_it.next().unwrap(); + let ub = i + 2 - i % 2; + assert_eq!(ub, k); + *v -= k; + } + + assert!(m_lower.mut_lower_bound(&199).next().is_none()); + + assert!(m_upper.mut_upper_bound(&198).next().is_none()); + + assert!(m_lower.iter().all(|(_, &x)| x == 0)); + assert!(m_upper.iter().all(|(_, &x)| x == 0)); + } + #[test] fn test_eq() { let mut a = TreeMap::new();