From 463a6d90643f45d9f86157e853ebd1a8402c6022 Mon Sep 17 00:00:00 2001 From: Huon Wilson Date: Mon, 6 Jan 2014 00:40:53 +1100 Subject: [PATCH] extra::treemap: add mutable-value iterators. --- src/libextra/treemap.rs | 279 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 278 insertions(+), 1 deletion(-) diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs index 1cf980b1059..49ee1ebebcc 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 @@ -157,6 +157,23 @@ 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 should be initialized using /// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`. fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> { @@ -212,6 +229,63 @@ impl TreeMap { } } } + /// 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 { @@ -340,6 +414,146 @@ fn iter_traverse_complete<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) { } } + +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> { + 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 +} + +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() + } +} + +/// 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 { priv stack: ~[TreeNode], @@ -1129,6 +1343,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();