//! 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
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 should be initialized using
/// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`.
fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
}
}
}
+ /// Get a lazy iterator that should be initialized using
+ /// `mut_iter_traverse_left`/`mut_iter_traverse_right`/`mut_iter_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> {
+ 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;
+ }
+ }
+ }
+
+ /// 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> {
+ 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<K, V> {
}
}
+
+fn mut_deref<K, V>(x: &mut Option<~TreeNode<K, V>>) -> *mut TreeNode<K, V> {
+ match *x {
+ Some(ref mut n) => {
+ let n: &mut TreeNode<K, V> = *n;
+ n as *mut TreeNode<K, V>
+ }
+ None => ptr::mut_null()
+ }
+}
+
+/// 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<K, V>],
+ // 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<K, V>,
+ priv remaining_min: uint,
+ 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<uint>) {
+ (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<uint>) {
+ self.iter.size_hint()
+ }
+}
+
+/// 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);
+}
+
+#[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);
+}
+
+/// 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);
+ }
+ it.node = ptr::mut_null();
+ }
+}
+
+
+
/// Lazy forward iterator over a map that consumes the map while iterating
pub struct TreeMapMoveIterator<K, V> {
priv stack: ~[TreeNode<K, V>],
}
}
+ #[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();
projects / rust.git / commitdiff