Skip to content

Commit 4db0611

Browse files
committed
auto merge of #8227 : dim-an/rust/tree-iter, r=thestinger
2 parents 7b2163d + 73ec9f3 commit 4db0611

File tree

1 file changed

+147
-5
lines changed

1 file changed

+147
-5
lines changed

src/libextra/treemap.rs

+147-5
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,68 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
184184
/// Get a lazy iterator over the key-value pairs in the map.
185185
/// Requires that it be frozen (immutable).
186186
pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
187-
TreeMapIterator{stack: ~[], node: &self.root, remaining: self.length}
187+
TreeMapIterator {
188+
stack: ~[],
189+
node: &self.root,
190+
remaining_min: self.length,
191+
remaining_max: self.length
192+
}
193+
}
194+
195+
/// Get a lazy iterator that should be initialized using
196+
/// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`.
197+
fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
198+
TreeMapIterator {
199+
stack: ~[],
200+
node: &self.root,
201+
remaining_min: 0,
202+
remaining_max: self.length
203+
}
204+
}
205+
206+
/// Return a lazy iterator to the first key-value pair whose key is not less than `k`
207+
/// If all keys in map are less than `k` an empty iterator is returned.
208+
pub fn lower_bound_iter<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> {
209+
let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal();
210+
loop {
211+
match *iter.node {
212+
Some(ref r) => {
213+
match k.cmp(&r.key) {
214+
Less => iter_traverse_left(&mut iter),
215+
Greater => iter_traverse_right(&mut iter),
216+
Equal => {
217+
iter_traverse_complete(&mut iter);
218+
return iter;
219+
}
220+
}
221+
}
222+
None => {
223+
iter_traverse_complete(&mut iter);
224+
return iter;
225+
}
226+
}
227+
}
228+
}
229+
230+
/// Return a lazy iterator to the first key-value pair whose key is greater than `k`
231+
/// If all keys in map are not greater than `k` an empty iterator is returned.
232+
pub fn upper_bound_iter<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> {
233+
let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal();
234+
loop {
235+
match *iter.node {
236+
Some(ref r) => {
237+
match k.cmp(&r.key) {
238+
Less => iter_traverse_left(&mut iter),
239+
Greater => iter_traverse_right(&mut iter),
240+
Equal => iter_traverse_right(&mut iter)
241+
}
242+
}
243+
None => {
244+
iter_traverse_complete(&mut iter);
245+
return iter;
246+
}
247+
}
248+
}
188249
}
189250

190251
/// Get a lazy iterator that consumes the treemap.
@@ -205,7 +266,8 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
205266
pub struct TreeMapIterator<'self, K, V> {
206267
priv stack: ~[&'self ~TreeNode<K, V>],
207268
priv node: &'self Option<~TreeNode<K, V>>,
208-
priv remaining: uint
269+
priv remaining_min: uint,
270+
priv remaining_max: uint
209271
}
210272

211273
impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V> {
@@ -222,7 +284,10 @@ impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V
222284
None => {
223285
let res = self.stack.pop();
224286
self.node = &res.right;
225-
self.remaining -= 1;
287+
self.remaining_max -= 1;
288+
if self.remaining_min > 0 {
289+
self.remaining_min -= 1;
290+
}
226291
return Some((&res.key, &res.value));
227292
}
228293
}
@@ -232,7 +297,46 @@ impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V
232297

233298
#[inline]
234299
fn size_hint(&self) -> (uint, Option<uint>) {
235-
(self.remaining, Some(self.remaining))
300+
(self.remaining_min, Some(self.remaining_max))
301+
}
302+
}
303+
304+
/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
305+
/// initialize TreeMapIterator pointing to element inside tree structure.
306+
///
307+
/// They should be used in following manner:
308+
/// - create iterator using TreeMap::iter_for_traversal
309+
/// - find required node using `iter_traverse_left`/`iter_traverse_right`
310+
/// (current node is `TreeMapIterator::node` field)
311+
/// - complete initialization with `iter_traverse_complete`
312+
#[inline]
313+
fn iter_traverse_left<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
314+
let node = it.node.get_ref();
315+
it.stack.push(node);
316+
it.node = &node.left;
317+
}
318+
319+
#[inline]
320+
fn iter_traverse_right<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
321+
it.node = &(it.node.get_ref().right);
322+
}
323+
324+
/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
325+
/// initialize TreeMapIterator pointing to element inside tree structure.
326+
///
327+
/// Completes traversal. Should be called before using iterator.
328+
/// Iteration will start from `self.node`.
329+
/// If `self.node` is None iteration will start from last node from which we
330+
/// traversed left.
331+
#[inline]
332+
fn iter_traverse_complete<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
333+
static none: Option<~TreeNode<K, V>> = None;
334+
match *it.node {
335+
Some(ref n) => {
336+
it.stack.push(n);
337+
it.node = &none;
338+
}
339+
None => ()
236340
}
237341
}
238342

@@ -417,6 +521,20 @@ impl<T: TotalOrd> TreeSet<T> {
417521
TreeSetIterator{iter: self.map.iter()}
418522
}
419523

524+
/// Get a lazy iterator pointing to the first value not less than `v` (greater or equal).
525+
/// If all elements in the set are less than `v` empty iterator is returned.
526+
#[inline]
527+
pub fn lower_bound_iter<'a>(&'a self, v: &T) -> TreeSetIterator<'a, T> {
528+
TreeSetIterator{iter: self.map.lower_bound_iter(v)}
529+
}
530+
531+
/// Get a lazy iterator pointing to the first value greater than `v`.
532+
/// If all elements in the set are not greater than `v` empty iterator is returned.
533+
#[inline]
534+
pub fn upper_bound_iter<'a>(&'a self, v: &T) -> TreeSetIterator<'a, T> {
535+
TreeSetIterator{iter: self.map.upper_bound_iter(v)}
536+
}
537+
420538
/// Visit all values in reverse order
421539
#[inline]
422540
pub fn each_reverse(&self, f: &fn(&T) -> bool) -> bool {
@@ -983,6 +1101,31 @@ mod test_treemap {
9831101
assert_eq!(*v, n * 2);
9841102
n += 1;
9851103
}
1104+
assert_eq!(n, 5);
1105+
}
1106+
1107+
#[test]
1108+
fn test_interval_iteration() {
1109+
let mut m = TreeMap::new();
1110+
for i in range(1, 100) {
1111+
assert!(m.insert(i * 2, i * 4));
1112+
}
1113+
1114+
for i in range(1, 198) {
1115+
let mut lb_it = m.lower_bound_iter(&i);
1116+
let (&k, &v) = lb_it.next().unwrap();
1117+
let lb = i + i % 2;
1118+
assert_eq!(lb, k);
1119+
assert_eq!(lb * 2, v);
1120+
1121+
let mut ub_it = m.upper_bound_iter(&i);
1122+
let (&k, &v) = ub_it.next().unwrap();
1123+
let ub = i + 2 - i % 2;
1124+
assert_eq!(ub, k);
1125+
assert_eq!(ub * 2, v);
1126+
}
1127+
let mut end_it = m.lower_bound_iter(&199);
1128+
assert_eq!(end_it.next(), None);
9861129
}
9871130

9881131
#[test]
@@ -1256,7 +1399,6 @@ mod test_set {
12561399

12571400
let mut n = 0;
12581401
for x in m.iter() {
1259-
printfln!(x);
12601402
assert_eq!(*x, n);
12611403
n += 1
12621404
}

0 commit comments

Comments
 (0)