extra: add consume iter to treemap.

This commit is contained in:
Graydon Hoare 2013-07-18 12:37:18 -07:00
parent 31c180e5f5
commit 9e4ebdb9d6

View File

@ -204,6 +204,19 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> { pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
TreeMapIterator{stack: ~[], node: &self.root, remaining: self.length} TreeMapIterator{stack: ~[], node: &self.root, remaining: self.length}
} }
/// Get a lazy iterator that consumes the treemap.
pub fn consume_iter(self) -> TreeMapConsumeIterator<K, V> {
let TreeMap { root: root, length: length } = self;
let stk = match root {
None => ~[],
Some(~tn) => ~[tn]
};
TreeMapConsumeIterator {
stack: stk,
remaining: length
}
}
} }
/// Lazy forward iterator over a map /// Lazy forward iterator over a map
@ -241,6 +254,56 @@ impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V
} }
} }
/// Lazy forward iterator over a map that consumes the map while iterating
pub struct TreeMapConsumeIterator<K, V> {
priv stack: ~[TreeNode<K, V>],
priv remaining: uint
}
impl<K, V> Iterator<(K, V)> for TreeMapConsumeIterator<K,V> {
#[inline]
fn next(&mut self) -> Option<(K, V)> {
while !self.stack.is_empty() {
let TreeNode {
key: key,
value: value,
left: left,
right: right,
level: level
} = self.stack.pop();
match left {
Some(~left) => {
let n = TreeNode {
key: key,
value: value,
left: None,
right: right,
level: level
};
self.stack.push(n);
self.stack.push(left);
}
None => {
match right {
Some(~right) => self.stack.push(right),
None => ()
}
self.remaining -= 1;
return Some((key, value))
}
}
}
None
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
(self.remaining, Some(self.remaining))
}
}
impl<'self, T> Iterator<&'self T> for TreeSetIterator<'self, T> { impl<'self, T> Iterator<&'self T> for TreeSetIterator<'self, T> {
/// Advance the iterator to the next node (in order). If there are no more nodes, return `None`. /// Advance the iterator to the next node (in order). If there are no more nodes, return `None`.
#[inline] #[inline]