Skip to content

Commit 8e492cc

Browse files
borsbors
committed
auto merge of #5123 : thestinger/rust/treemap, r=nikomatsakis
* replace the dual next() and get() calls with a single next() function * drop one of the pointer members from the struct * add a method for using the lazy iterator with a for loop
2 parents f4e8ac2 + 5b0a2d1 commit 8e492cc

File tree

1 file changed

+93
-113
lines changed

1 file changed

+93
-113
lines changed

‎src/libstd/treemap.rs

+93-113
Original file line numberDiff line numberDiff line change
@@ -48,13 +48,8 @@ impl<K:Eq + Ord,V:Eq> Eq for TreeMap<K, V> {
4848
let mut y = other.iter();
4949
for self.len().times {
5050
unsafe { // unsafe as a purity workaround
51-
map_next(&mut x);
52-
map_next(&mut y);
53-
// FIXME: #4492 (ICE), x.get() == y.get()
54-
let (x1, x2) = x.get().unwrap();
55-
let (y1, y2) = y.get().unwrap();
56-
57-
if x1 != y1 || x2 != y2 {
51+
if map_next(&mut x).unwrap() !=
52+
map_next(&mut y).unwrap() {
5853
return false
5954
}
6055
}
@@ -73,10 +68,8 @@ pure fn lt<K:Ord,V>(a: &TreeMap<K, V>, b: &TreeMap<K, V>) -> bool {
7368
let (a_len, b_len) = (a.len(), b.len());
7469
for uint::min(a_len, b_len).times {
7570
unsafe { // purity workaround
76-
map_next(&mut x);
77-
map_next(&mut y);
78-
let (key_a,_) = x.get().unwrap();
79-
let (key_b,_) = y.get().unwrap();
71+
let (key_a,_) = map_next(&mut x).unwrap();
72+
let (key_b,_) = map_next(&mut y).unwrap();
8073
if *key_a < *key_b { return true; }
8174
if *key_a > *key_b { return false; }
8275
}
@@ -201,30 +194,21 @@ impl <K:Ord,V> TreeMap<K, V> {
201194
/// Get a lazy iterator over the key-value pairs in the map.
202195
/// Requires that it be frozen (immutable).
203196
pure fn iter(&self) -> TreeMapIterator/&self<K, V> {
204-
TreeMapIterator{stack: ~[], node: &self.root, current: None}
197+
TreeMapIterator{stack: ~[], node: &self.root}
205198
}
206199
}
207200

208201
/// Lazy forward iterator over a map
209202
pub struct TreeMapIterator<K, V> {
210203
priv stack: ~[&~TreeNode<K, V>],
211-
priv node: &Option<~TreeNode<K, V>>,
212-
priv current: Option<&~TreeNode<K, V>>
204+
priv node: &Option<~TreeNode<K, V>>
213205
}
214206

215-
impl <K:Ord,V> TreeMapIterator<K, V> {
216-
// Returns the current node, or None if this iterator is at the end.
217-
fn get(&const self) -> Option<(&self/K, &self/V)> {
218-
match self.current {
219-
Some(res) => Some((&res.key, &res.value)),
220-
None => None
221-
}
222-
}
223-
}
224-
225-
/// Advance the iterator to the next node (in order). If this iterator
226-
/// is finished, does nothing.
227-
pub fn map_next<K:Ord,V>(iter: &mut TreeMapIterator/&a<K, V>) {
207+
/// Advance the iterator to the next node (in order) and return a
208+
/// tuple with a reference to the key and value. If there are no
209+
/// more nodes, return `None`.
210+
fn map_next<K: Ord, V>(iter: &mut TreeMapIterator/&r<K, V>)
211+
-> Option<(&r/K, &r/V)> {
228212
while !iter.stack.is_empty() || iter.node.is_some() {
229213
match *iter.node {
230214
Some(ref x) => {
@@ -234,12 +218,24 @@ pub fn map_next<K:Ord,V>(iter: &mut TreeMapIterator/&a<K, V>) {
234218
None => {
235219
let res = iter.stack.pop();
236220
iter.node = &res.right;
237-
iter.current = Some(res);
238-
return;
221+
return Some((&res.key, &res.value));
239222
}
240223
}
241224
}
242-
iter.current = None;
225+
None
226+
}
227+
228+
/// Advance the iterator through the map
229+
fn map_advance<K: Ord, V>(iter: &mut TreeMapIterator/&r<K, V>,
230+
f: fn((&r/K, &r/V)) -> bool) {
231+
loop {
232+
match map_next(iter) {
233+
Some(x) => {
234+
if !f(x) { return }
235+
}
236+
None => return
237+
}
238+
}
243239
}
244240

245241
pub struct TreeSet<T> {
@@ -308,19 +304,15 @@ impl<T:Ord> Set<T> for TreeSet<T> {
308304
let mut x = self.iter();
309305
let mut y = other.iter();
310306
unsafe { // purity workaround
311-
set_next(&mut x);
312-
set_next(&mut y);
313-
let mut a = x.get();
314-
let mut b = y.get();
307+
let mut a = set_next(&mut x);
308+
let mut b = set_next(&mut y);
315309
while a.is_some() && b.is_some() {
316310
let a1 = a.unwrap();
317311
let b1 = b.unwrap();
318312
if a1 < b1 {
319-
set_next(&mut x);
320-
a = x.get();
313+
a = set_next(&mut x);
321314
} else if b1 < a1 {
322-
set_next(&mut y);
323-
b = y.get();
315+
b = set_next(&mut y);
324316
} else {
325317
return false;
326318
}
@@ -339,10 +331,8 @@ impl<T:Ord> Set<T> for TreeSet<T> {
339331
let mut x = self.iter();
340332
let mut y = other.iter();
341333
unsafe { // purity workaround
342-
set_next(&mut x);
343-
set_next(&mut y);
344-
let mut a = x.get();
345-
let mut b = y.get();
334+
let mut a = set_next(&mut x);
335+
let mut b = set_next(&mut y);
346336
while b.is_some() {
347337
if a.is_none() {
348338
return false
@@ -356,11 +346,9 @@ impl<T:Ord> Set<T> for TreeSet<T> {
356346
}
357347

358348
if !(a1 < b1) {
359-
set_next(&mut y);
360-
b = y.get();
349+
b = set_next(&mut y);
361350
}
362-
set_next(&mut x);
363-
a = x.get();
351+
a = set_next(&mut x);
364352
}
365353
}
366354
true
@@ -372,15 +360,13 @@ impl<T:Ord> Set<T> for TreeSet<T> {
372360
let mut y = other.iter();
373361

374362
unsafe { // purity workaround
375-
set_next(&mut x);
376-
set_next(&mut y);
377-
let mut a = x.get();
378-
let mut b = y.get();
363+
let mut a = set_next(&mut x);
364+
let mut b = set_next(&mut y);
379365

380366
while a.is_some() {
381367
if b.is_none() {
382368
return do a.while_some() |a1| {
383-
if f(a1) { set_next(&mut x); x.get() } else { None }
369+
if f(a1) { set_next(&mut x) } else { None }
384370
}
385371
}
386372

@@ -389,12 +375,10 @@ impl<T:Ord> Set<T> for TreeSet<T> {
389375

390376
if a1 < b1 {
391377
if !f(a1) { return }
392-
set_next(&mut x);
393-
a = x.get();
378+
a = set_next(&mut x);
394379
} else {
395-
if !(b1 < a1) { set_next(&mut x); a = x.get() }
396-
set_next(&mut y);
397-
b = y.get();
380+
if !(b1 < a1) { a = set_next(&mut x) }
381+
b = set_next(&mut y);
398382
}
399383
}
400384
}
@@ -407,15 +391,13 @@ impl<T:Ord> Set<T> for TreeSet<T> {
407391
let mut y = other.iter();
408392

409393
unsafe { // purity workaround
410-
set_next(&mut x);
411-
set_next(&mut y);
412-
let mut a = x.get();
413-
let mut b = y.get();
394+
let mut a = set_next(&mut x);
395+
let mut b = set_next(&mut y);
414396

415397
while a.is_some() {
416398
if b.is_none() {
417399
return do a.while_some() |a1| {
418-
if f(a1) { set_next(&mut x); x.get() } else { None }
400+
if f(a1) { set_next(&mut x) } else { None }
419401
}
420402
}
421403

@@ -424,21 +406,18 @@ impl<T:Ord> Set<T> for TreeSet<T> {
424406

425407
if a1 < b1 {
426408
if !f(a1) { return }
427-
set_next(&mut x);
428-
a = x.get();
409+
a = set_next(&mut x);
429410
} else {
430411
if b1 < a1 {
431412
if !f(b1) { return }
432413
} else {
433-
set_next(&mut x);
434-
a = x.get();
414+
a = set_next(&mut x);
435415
}
436-
set_next(&mut y);
437-
b = y.get();
416+
b = set_next(&mut y);
438417
}
439418
}
440419
do b.while_some |b1| {
441-
if f(b1) { set_next(&mut y); y.get() } else { None }
420+
if f(b1) { set_next(&mut y) } else { None }
442421
}
443422
}
444423
}
@@ -449,23 +428,19 @@ impl<T:Ord> Set<T> for TreeSet<T> {
449428
let mut y = other.iter();
450429

451430
unsafe { // purity workaround
452-
set_next(&mut x);
453-
set_next(&mut y);
454-
let mut a = x.get();
455-
let mut b = y.get();
431+
let mut a = set_next(&mut x);
432+
let mut b = set_next(&mut y);
456433

457434
while a.is_some() && b.is_some() {
458435
let a1 = a.unwrap();
459436
let b1 = b.unwrap();
460437
if a1 < b1 {
461-
set_next(&mut x);
462-
a = x.get();
438+
a = set_next(&mut x);
463439
} else {
464440
if !(b1 < a1) {
465441
if !f(a1) { return }
466442
}
467-
set_next(&mut y);
468-
b = y.get();
443+
b = set_next(&mut y);
469444
}
470445
}
471446
}
@@ -477,15 +452,13 @@ impl<T:Ord> Set<T> for TreeSet<T> {
477452
let mut y = other.iter();
478453

479454
unsafe { // purity workaround
480-
set_next(&mut x);
481-
set_next(&mut y);
482-
let mut a = x.get();
483-
let mut b = y.get();
455+
let mut a = set_next(&mut x);
456+
let mut b = set_next(&mut y);
484457

485458
while a.is_some() {
486459
if b.is_none() {
487460
return do a.while_some() |a1| {
488-
if f(a1) { set_next(&mut x); x.get() } else { None }
461+
if f(a1) { set_next(&mut x) } else { None }
489462
}
490463
}
491464

@@ -494,16 +467,13 @@ impl<T:Ord> Set<T> for TreeSet<T> {
494467

495468
if b1 < a1 {
496469
if !f(b1) { return }
497-
set_next(&mut y);
498-
b = y.get();
470+
b = set_next(&mut y);
499471
} else {
500472
if !f(a1) { return }
501473
if !(a1 < b1) {
502-
set_next(&mut y);
503-
b = y.get()
474+
b = set_next(&mut y);
504475
}
505-
set_next(&mut x);
506-
a = x.get();
476+
a = set_next(&mut x);
507477
}
508478
}
509479
}
@@ -526,20 +496,16 @@ pub struct TreeSetIterator<T> {
526496
priv iter: TreeMapIterator<T, ()>
527497
}
528498

529-
impl <T:Ord> TreeSetIterator<T> {
530-
/// Returns the current node, or None if this iterator is at the end.
531-
fn get(&const self) -> Option<&self/T> {
532-
match self.iter.get() {
533-
None => None,
534-
Some((k, _)) => Some(k)
535-
}
536-
}
537-
}
538-
539499
/// Advance the iterator to the next node (in order). If this iterator is
540500
/// finished, does nothing.
541-
pub fn set_next<T:Ord>(iter: &mut TreeSetIterator/&a<T>) {
542-
map_next(&mut iter.iter);
501+
pub fn set_next<T: Ord>(iter: &mut TreeSetIterator/&r<T>) -> Option<&r/T> {
502+
do map_next(&mut iter.iter).map |&(value, _)| { value }
503+
}
504+
505+
/// Advance the iterator through the set
506+
fn set_advance<T: Ord>(iter: &mut TreeSetIterator/&r<T>,
507+
f: fn(&r/T) -> bool) {
508+
do map_advance(&mut iter.iter) |(k, _)| { f(k) }
543509
}
544510

545511
// Nodes keep track of their level in the tree, starting at 1 in the
@@ -983,23 +949,37 @@ mod test_treemap {
983949
assert m.insert(x5, y5);
984950

985951
let m = m;
986-
let mut iter = m.iter();
952+
let mut a = m.iter();
987953

988954
// FIXME: #4492 (ICE): iter.get() == Some((&x1, &y1))
989955

990-
map_next(&mut iter);
991-
assert iter.get().unwrap() == (&x1, &y1);
992-
map_next(&mut iter);
993-
assert iter.get().unwrap() == (&x2, &y2);
994-
map_next(&mut iter);
995-
assert iter.get().unwrap() == (&x3, &y3);
996-
map_next(&mut iter);
997-
assert iter.get().unwrap() == (&x4, &y4);
998-
map_next(&mut iter);
999-
assert iter.get().unwrap() == (&x5, &y5);
1000-
1001-
map_next(&mut iter);
1002-
assert iter.get().is_none();
956+
assert map_next(&mut a).unwrap() == (&x1, &y1);
957+
assert map_next(&mut a).unwrap() == (&x2, &y2);
958+
assert map_next(&mut a).unwrap() == (&x3, &y3);
959+
assert map_next(&mut a).unwrap() == (&x4, &y4);
960+
assert map_next(&mut a).unwrap() == (&x5, &y5);
961+
962+
assert map_next(&mut a).is_none();
963+
964+
let mut b = m.iter();
965+
966+
let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4),
967+
(&x5, &y5)];
968+
let mut i = 0;
969+
970+
for map_advance(&mut b) |x| {
971+
assert expected[i] == x;
972+
i += 1;
973+
974+
if i == 2 {
975+
break
976+
}
977+
}
978+
979+
for map_advance(&mut b) |x| {
980+
assert expected[i] == x;
981+
i += 1;
982+
}
1003983
}
1004984
}
1005985

0 commit comments

Comments
 (0)