extra: Implement .rev_iter() in treemap

Implement reverse iterators for TreeMap and TreeSet, that produce the
keys in backward order.

diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs
index 051587a74f963262949dbd80927afeb33630f02e..989f7a419bba63d0ead0731378b66b2482b3b80f 100644 (file)
--- a/src/libextra/treemap.rs
+++ b/src/libextra/treemap.rs
@@ -13,7 +13,6 @@
 //! `TotalOrd`.
 
 
-use std::num;
 use std::util::{swap, replace};
 use std::iterator::{FromIterator, Extendable};
 
@@ -152,7 +151,7 @@ pub fn mutate_values(&mut self, f: &fn(&K, &mut V) -> bool) -> bool {
 
     /// Visit all key-value pairs in reverse order
     pub fn each_reverse<'a>(&'a self, f: &fn(&'a K, &'a V) -> bool) -> bool {
-        each_reverse(&self.root, f)
+        self.rev_iter().advance(|(k,v)| f(k, v))
     }
 
     /// Visit all keys in reverse order
@@ -176,6 +175,12 @@ pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
         }
     }
 
+    /// Get a lazy reverse iterator over the key-value pairs in the map.
+    /// Requires that it be frozen (immutable).
+    pub fn rev_iter<'a>(&'a self) -> TreeMapRevIterator<'a, K, V> {
+        TreeMapRevIterator{iter: self.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> {
@@ -254,20 +259,18 @@ pub struct TreeMapIterator<'self, K, V> {
     priv remaining_max: uint
 }
 
-impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, 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<(&'self K, &'self V)> {
+impl<'self, K, V> TreeMapIterator<'self, K, V> {
+    #[inline(always)]
+    fn next_(&mut self, forward: bool) -> Option<(&'self K, &'self V)> {
         while !self.stack.is_empty() || self.node.is_some() {
             match *self.node {
               Some(ref x) => {
                 self.stack.push(x);
-                self.node = &x.left;
+                self.node = if forward { &x.left } else { &x.right };
               }
               None => {
                 let res = self.stack.pop();
-                self.node = &res.right;
+                self.node = if forward { &res.right } else { &res.left };
                 self.remaining_max -= 1;
                 if self.remaining_min > 0 {
                     self.remaining_min -= 1;
@@ -278,6 +281,15 @@ fn next(&mut self) -> Option<(&'self K, &'self V)> {
         }
         None
     }
+}
+
+impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, 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<(&'self K, &'self V)> {
+        self.next_(true)
+    }
 
     #[inline]
     fn size_hint(&self) -> (uint, Option<uint>) {
@@ -285,6 +297,25 @@ fn size_hint(&self) -> (uint, Option<uint>) {
     }
 }
 
+/// Lazy backward iterator over a map
+pub struct TreeMapRevIterator<'self, K, V> {
+    priv iter: TreeMapIterator<'self, K, V>,
+}
+
+impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapRevIterator<'self, 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<(&'self K, &'self 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 TreeMapIterator pointing to element inside tree structure.
 ///
@@ -382,6 +413,14 @@ fn next(&mut self) -> Option<&'self T> {
     }
 }
 
+impl<'self, T> Iterator<&'self T> for TreeSetRevIterator<'self, T> {
+    /// Advance the iterator to the next node (in order). If there are no more nodes, return `None`.
+    #[inline]
+    fn next(&mut self) -> Option<&'self T> {
+        do self.iter.next().map |&(value, _)| { value }
+    }
+}
+
 /// A implementation of the `Set` trait on top of the `TreeMap` container. The
 /// only requirement is that the type of the elements contained ascribes to the
 /// `TotalOrd` trait.
@@ -492,6 +531,13 @@ pub fn iter<'a>(&'a self) -> TreeSetIterator<'a, T> {
         TreeSetIterator{iter: self.map.iter()}
     }
 
+    /// Get a lazy iterator over the values in the set.
+    /// Requires that it be frozen (immutable).
+    #[inline]
+    pub fn rev_iter<'a>(&'a self) -> TreeSetRevIterator<'a, T> {
+        TreeSetRevIterator{iter: self.map.rev_iter()}
+    }
+
     /// Get a lazy iterator pointing to the first value not less than `v` (greater or equal).
     /// If all elements in the set are less than `v` empty iterator is returned.
     #[inline]
@@ -540,6 +586,11 @@ pub struct TreeSetIterator<'self, T> {
     priv iter: TreeMapIterator<'self, T, ()>
 }
 
+/// Lazy backward iterator over a set
+pub struct TreeSetRevIterator<'self, T> {
+    priv iter: TreeMapRevIterator<'self, T, ()>
+}
+
 // Encapsulate an iterator and hold its latest value until stepped forward
 struct Focus<A, T> {
     priv iter: T,
@@ -691,18 +742,6 @@ pub fn new(key: K, value: V) -> TreeNode<K, V> {
     }
 }
 
-fn each<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>,
-                            f: &fn(&'r K, &'r V) -> bool) -> bool {
-    node.iter().advance(|x| each(&x.left,  |k,v| f(k,v)) && f(&x.key, &x.value) &&
-                            each(&x.right, |k,v| f(k,v)))
-}
-
-fn each_reverse<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>,
-                                    f: &fn(&'r K, &'r V) -> bool) -> bool {
-    node.iter().advance(|x| each_reverse(&x.right, |k,v| f(k,v)) && f(&x.key, &x.value) &&
-                            each_reverse(&x.left,  |k,v| f(k,v)))
-}
-
 fn mutate_values<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
                                      f: &fn(&'r K, &'r mut V) -> bool)
                                   -> bool {