]> git.lizzy.rs Git - rust.git/commitdiff
extra::treemap: share code between the mutable and immutable iterators.
authorHuon Wilson <dbau.pp+github@gmail.com>
Sun, 5 Jan 2014 14:53:27 +0000 (01:53 +1100)
committerHuon Wilson <dbau.pp+github@gmail.com>
Sun, 5 Jan 2014 14:53:27 +0000 (01:53 +1100)
Yay for macros.

src/libextra/treemap.rs

index 2182809f3aa733f940a2239de25cc9002fff05e4..f4fd81437fc6eaf4bad91ec78541d024564abfc5 100644 (file)
@@ -169,8 +169,57 @@ pub fn mut_rev_iter<'a>(&'a mut self) -> TreeMapMutRevIterator<'a, K, V> {
     }
 
 
+    /// Get a lazy iterator that consumes the treemap.
+    pub fn move_iter(self) -> TreeMapMoveIterator<K, V> {
+        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<K: TotalOrd, V> TreeMap<K, V> {
     /// 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 @@ fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
     /// 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 @@ fn mut_iter_for_traversal<'a>(&'a mut self) -> TreeMapMutIterator<'a, K, V> {
     /// 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 @@ pub fn mut_lower_bound<'a>(&'a mut self, k: &K) -> TreeMapMutIterator<'a, K, V>
     /// 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> {
-        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<K, V>],
-    priv node: Option<&'a TreeNode<K, V>>,
+    // 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<K,V>>).
+    priv node: *TreeNode<K, V>,
     priv remaining_min: uint,
     priv remaining_max: uint
 }
 
-fn deref<'a, K, V>(node: &'a Option<~TreeNode<K, V>>) -> Option<&'a TreeNode<K, V>> {
-    match *node {
-        Some(ref n) => {
-            let n: &TreeNode<K, V> = *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<uint>) {
-        (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<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.
-///
-/// 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<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> {
@@ -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);
+/// Lazy backward iterator over a map
+pub struct TreeMapMutRevIterator<'a, K, V> {
+    priv iter: TreeMapMutIterator<'a, K, V>,
+}
+
+
+// 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 `& <mutability> <operand>` token
+// sequence and forces their evalutation as an expression.
+macro_rules! addr { ($e:expr) => { $e }}
+
+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,
+
+     // the function to go from &m Option<~TreeNode> to *m TreeNode
+     deref = $deref:ident,
+
+     // 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);
-            } 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;
+            }
+
+            #[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();
                 }
-                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)
-    }
+        // 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<uint>) {
-        (self.remaining_min, Some(self.remaining_max))
-    }
-}
+            #[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>,
-}
+        // 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)
+            }
 
-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()
+            }
+        }
     }
+} // end of define_iterator
 
-    #[inline]
-    fn size_hint(&self) -> (uint, Option<uint>) {
-        self.iter.size_hint()
-    }
+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,
 
-/// 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);
+    addr_mut = mut
 }
 
-#[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);
+fn deref<'a, K, V>(node: &'a Option<~TreeNode<K, V>>) -> *TreeNode<K, V> {
+    match *node {
+        Some(ref n) => {
+            let n: &TreeNode<K, V> = *n;
+            n as *TreeNode<K, V>
+        }
+        None => ptr::null()
+    }
 }
 
-/// 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);
+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>
         }
-        it.node = ptr::mut_null();
+        None => ptr::mut_null()
     }
 }