extra::treemap: add mutable-value iterators.

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<K: TotalOrd, V> TreeMap<K, V> {
         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<K: TotalOrd, V> TreeMap<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> {
@@ -340,6 +414,146 @@ fn iter_traverse_complete<'a, K, V>(it: &mut TreeMapIterator<'a, 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>],
@@ -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();