diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs index 2182809f3aa..f4fd81437fc 100644 --- a/src/libextra/treemap.rs +++ b/src/libextra/treemap.rs @@ -169,8 +169,57 @@ impl TreeMap { } + /// 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: ~[], @@ -183,49 +232,17 @@ 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 should be initialized using - /// `mut_iter_traverse_left`/`mut_iter_traverse_right`/`mut_iter_traverse_complete`. + /// `traverse_left`/`traverse_right`/`traverse_complete`. fn mut_iter_for_traversal<'a>(&'a mut self) -> TreeMapMutIterator<'a, K, V> { TreeMapMutIterator { stack: ~[], @@ -241,23 +258,7 @@ impl TreeMap { /// 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> { - let mut iter = self.mut_iter_for_traversal(); - loop { - if !iter.node.is_null() { - let node_k = unsafe {&(*iter.node).key}; - match k.cmp(node_k) { - Less => mut_iter_traverse_left(&mut iter), - Greater => mut_iter_traverse_right(&mut iter), - Equal => { - mut_iter_traverse_complete(&mut iter); - return iter; - } - } - } else { - mut_iter_traverse_complete(&mut iter); - return iter; - } - } + bound_setup!(self.mut_iter_for_traversal(), true) } /// Return a lazy iterator to the first key-value pair (with the @@ -266,160 +267,26 @@ impl TreeMap { /// 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> { - let mut iter = self.mut_iter_for_traversal(); - loop { - if !iter.node.is_null() { - let node_k = unsafe {&(*iter.node).key}; - match k.cmp(node_k) { - Less => mut_iter_traverse_left(&mut iter), - Greater => mut_iter_traverse_right(&mut iter), - Equal => mut_iter_traverse_right(&mut iter) - } - } else { - mut_iter_traverse_complete(&mut iter); - return 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 - } + 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 - } -} - -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; - } - 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) - } - - #[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>, } -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() - } -} - -/// 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); -} - -#[inline] -fn iter_traverse_right<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) { - it.node = deref(&it.node.get_ref().right); -} - -/// 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; - } - 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 => ptr::mut_null() - } -} - /// Lazy forward iterator over a map that allows for the mutation of /// the values. pub struct TreeMapMutIterator<'a, K, V> { @@ -448,102 +315,169 @@ pub struct TreeMapMutIterator<'a, K, V> { priv remaining_max: uint } -impl<'a, K, V> TreeMapMutIterator<'a, K, V> { - #[inline(always)] - fn next_(&mut self, forward: bool) -> Option<(&'a K, &'a mut V)> { - while !self.stack.is_empty() || !self.node.is_null() { - let node = self.node; - if !node.is_null() { - let node = unsafe {&mut *node}; - { - let next_node = if forward { &mut node.left } else { &mut node.right }; - self.node = mut_deref(next_node); - } - self.stack.push(node); - } else { - let node = self.stack.pop(); - self.node = mut_deref(if forward { &mut node.right } else { &mut node.left }); - self.remaining_max -= 1; - if self.remaining_min > 0 { - self.remaining_min -= 1; - } - return Some((&node.key, &mut node.value)); - } - } - None - } -} - -impl<'a, K, V> Iterator<(&'a K, &'a mut V)> for TreeMapMutIterator<'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 mut V)> { - self.next_(true) - } - - #[inline] - fn size_hint(&self) -> (uint, Option) { - (self.remaining_min, Some(self.remaining_max)) - } -} - /// Lazy backward iterator over a map pub struct TreeMapMutRevIterator<'a, K, V> { priv iter: TreeMapMutIterator<'a, K, V>, } -impl<'a, K, V> Iterator<(&'a K, &'a mut V)> for TreeMapMutRevIterator<'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 mut V)> { - self.iter.next_(false) - } - #[inline] - fn size_hint(&self) -> (uint, Option) { - self.iter.size_hint() - } -} +// 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 }} -/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to -/// initialize TreeMapMutIterator 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 `TreeMapMutIterator::node` field) -/// - complete initialization with `iter_traverse_complete` -#[inline] -fn mut_iter_traverse_left<'a, K, V>(it: &mut TreeMapMutIterator<'a, K, V>) { - // guaranteed to be non-null - let node = unsafe {&mut *it.node}; - it.node = mut_deref(&mut node.left); - it.stack.push(node); -} +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, -#[inline] -fn mut_iter_traverse_right<'a, K, V>(it: &mut TreeMapMutIterator<'a, K, V>) { - // guaranteed to be non-null - let node = unsafe {&mut *it.node}; - it.node = mut_deref(&mut node.right); -} + // the function to go from &m Option<~TreeNode> to *m TreeNode + deref = $deref:ident, -/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to -/// initialize TreeMapMutIterator 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 mut_iter_traverse_complete<'a, K, V>(it: &mut TreeMapMutIterator<'a, K, V>) { - if !it.node.is_null() { - unsafe { - it.stack.push(&mut *it.node); + // 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(); + } + } } - it.node = ptr::mut_null(); + + // 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)) + } + } + + // 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) + } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + self.iter.size_hint() + } + } + } +} // end of define_iterator + +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, + + addr_mut = mut +} + +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() + } +} + +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 => ptr::mut_null() } }