1 // Copyright 2013 The Rust Project Developers. See the COPYRIGHT
2 // file at the top-level directory of this distribution and at
3 // http://rust-lang.org/COPYRIGHT.
4 //
5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8 // option. This file may not be copied, modified, or distributed
9 // except according to those terms.
10
11 //! An ordered map and set implemented as self-balancing binary search
12 //! trees. The only requirement for the types is that the key implements
13 //! `TotalOrd`.
14
15 use std::iter::{Peekable};
16 use std::cmp::Ordering;
17 use std::mem::{replace, swap};
18 use std::ptr;
19
20 // This is implemented as an AA tree, which is a simplified variation of
21 // a red-black tree where red (horizontal) nodes can only be added
22 // as a right child. The time complexity is the same, and re-balancing
23 // operations are more frequent but also cheaper.
24
25 // Future improvements:
26
27 // range search - O(log n) retrieval of an iterator from some key
28
29 // (possibly) implement the overloads Python does for sets:
30 // * intersection: &
31 // * difference: -
32 // * symmetric difference: ^
33 // * union: |
34 // These would be convenient since the methods work like `each`
35
36 #[allow(missing_doc)]
37 #[deriving(Clone)]
38 pub struct TreeMap<K, V> {
39 root: Option<Box<TreeNode<K, V>>>,
40 length: uint
41 }
42
43 impl<K: Eq + TotalOrd, V: Eq> Eq for TreeMap<K, V> {
44 fn eq(&self, other: &TreeMap<K, V>) -> bool {
45 self.len() == other.len() &&
46 self.iter().zip(other.iter()).all(|(a, b)| a == b)
47 }
48 }
49
50 // Lexicographical comparison
51 fn lt<K: Ord + TotalOrd, V: Ord>(a: &TreeMap<K, V>,
52 b: &TreeMap<K, V>) -> bool {
53 // the Zip iterator is as long as the shortest of a and b.
54 for ((key_a, value_a), (key_b, value_b)) in a.iter().zip(b.iter()) {
55 if *key_a < *key_b { return true; }
56 if *key_a > *key_b { return false; }
57 if *value_a < *value_b { return true; }
58 if *value_a > *value_b { return false; }
59 }
60
61 a.len() < b.len()
62 }
63
64 impl<K: Ord + TotalOrd, V: Ord> Ord for TreeMap<K, V> {
65 #[inline]
66 fn lt(&self, other: &TreeMap<K, V>) -> bool { lt(self, other) }
67 }
68
69 impl<K: TotalOrd, V> Container for TreeMap<K, V> {
70 fn len(&self) -> uint { self.length }
71 }
72
73 impl<K: TotalOrd, V> Mutable for TreeMap<K, V> {
74 fn clear(&mut self) {
75 self.root = None;
76 self.length = 0
77 }
78 }
79
80 impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
81 fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
82 let mut current: &'a Option<Box<TreeNode<K, V>>> = &self.root;
83 loop {
84 match *current {
85 Some(ref r) => {
86 match key.cmp(&r.key) {
87 Less => current = &r.left,
88 Greater => current = &r.right,
89 Equal => return Some(&r.value)
90 }
91 }
92 None => return None
93 }
94 }
95 }
96 }
97
98 impl<K: TotalOrd, V> MutableMap<K, V> for TreeMap<K, V> {
99 #[inline]
100 fn find_mut<'a>(&'a mut self, key: &K) -> Option<&'a mut V> {
101 find_mut(&mut self.root, key)
102 }
103
104 fn swap(&mut self, key: K, value: V) -> Option<V> {
105 let ret = insert(&mut self.root, key, value);
106 if ret.is_none() { self.length += 1 }
107 ret
108 }
109
110 fn pop(&mut self, key: &K) -> Option<V> {
111 let ret = remove(&mut self.root, key);
112 if ret.is_some() { self.length -= 1 }
113 ret
114 }
115 }
116
117 impl<K: TotalOrd, V> TreeMap<K, V> {
118 /// Create an empty TreeMap
119 pub fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} }
120
121 /// Get a lazy iterator over the key-value pairs in the map.
122 /// Requires that it be frozen (immutable).
123 pub fn iter<'a>(&'a self) -> Entries<'a, K, V> {
124 Entries {
125 stack: vec!(),
126 node: deref(&self.root),
127 remaining_min: self.length,
128 remaining_max: self.length
129 }
130 }
131
132 /// Get a lazy reverse iterator over the key-value pairs in the map.
133 /// Requires that it be frozen (immutable).
134 pub fn rev_iter<'a>(&'a self) -> RevEntries<'a, K, V> {
135 RevEntries{iter: self.iter()}
136 }
137
138 /// Get a lazy forward iterator over the key-value pairs in the
139 /// map, with the values being mutable.
140 pub fn mut_iter<'a>(&'a mut self) -> MutEntries<'a, K, V> {
141 MutEntries {
142 stack: vec!(),
143 node: mut_deref(&mut self.root),
144 remaining_min: self.length,
145 remaining_max: self.length
146 }
147 }
148 /// Get a lazy reverse iterator over the key-value pairs in the
149 /// map, with the values being mutable.
150 pub fn mut_rev_iter<'a>(&'a mut self) -> RevMutEntries<'a, K, V> {
151 RevMutEntries{iter: self.mut_iter()}
152 }
153
154
155 /// Get a lazy iterator that consumes the treemap.
156 pub fn move_iter(self) -> MoveEntries<K, V> {
157 let TreeMap { root: root, length: length } = self;
158 let stk = match root {
159 None => vec!(),
160 Some(box tn) => vec!(tn)
161 };
162 MoveEntries {
163 stack: stk,
164 remaining: length
165 }
166 }
167 }
168
169 // range iterators.
170
171 macro_rules! bound_setup {
172 // initialiser of the iterator to manipulate
173 ($iter:expr,
174 // whether we are looking for the lower or upper bound.
175 $is_lower_bound:expr) => {
176 {
177 let mut iter = $iter;
178 loop {
179 if !iter.node.is_null() {
180 let node_k = unsafe {&(*iter.node).key};
181 match k.cmp(node_k) {
182 Less => iter.traverse_left(),
183 Greater => iter.traverse_right(),
184 Equal => {
185 if $is_lower_bound {
186 iter.traverse_complete();
187 return iter;
188 } else {
189 iter.traverse_right()
190 }
191 }
192 }
193 } else {
194 iter.traverse_complete();
195 return iter;
196 }
197 }
198 }
199 }
200 }
201
202
203 impl<K: TotalOrd, V> TreeMap<K, V> {
204 /// Get a lazy iterator that should be initialized using
205 /// `traverse_left`/`traverse_right`/`traverse_complete`.
206 fn iter_for_traversal<'a>(&'a self) -> Entries<'a, K, V> {
207 Entries {
208 stack: vec!(),
209 node: deref(&self.root),
210 remaining_min: 0,
211 remaining_max: self.length
212 }
213 }
214
215 /// Return a lazy iterator to the first key-value pair whose key is not less than `k`
216 /// If all keys in map are less than `k` an empty iterator is returned.
217 pub fn lower_bound<'a>(&'a self, k: &K) -> Entries<'a, K, V> {
218 bound_setup!(self.iter_for_traversal(), true)
219 }
220
221 /// Return a lazy iterator to the first key-value pair whose key is greater than `k`
222 /// If all keys in map are not greater than `k` an empty iterator is returned.
223 pub fn upper_bound<'a>(&'a self, k: &K) -> Entries<'a, K, V> {
224 bound_setup!(self.iter_for_traversal(), false)
225 }
226
227 /// Get a lazy iterator that should be initialized using
228 /// `traverse_left`/`traverse_right`/`traverse_complete`.
229 fn mut_iter_for_traversal<'a>(&'a mut self) -> MutEntries<'a, K, V> {
230 MutEntries {
231 stack: vec!(),
232 node: mut_deref(&mut self.root),
233 remaining_min: 0,
234 remaining_max: self.length
235 }
236 }
237
238 /// Return a lazy value iterator to the first key-value pair (with
239 /// the value being mutable) whose key is not less than `k`.
240 ///
241 /// If all keys in map are less than `k` an empty iterator is
242 /// returned.
243 pub fn mut_lower_bound<'a>(&'a mut self, k: &K) -> MutEntries<'a, K, V> {
244 bound_setup!(self.mut_iter_for_traversal(), true)
245 }
246
247 /// Return a lazy iterator to the first key-value pair (with the
248 /// value being mutable) whose key is greater than `k`.
249 ///
250 /// If all keys in map are not greater than `k` an empty iterator
251 /// is returned.
252 pub fn mut_upper_bound<'a>(&'a mut self, k: &K) -> MutEntries<'a, K, V> {
253 bound_setup!(self.mut_iter_for_traversal(), false)
254 }
255 }
256
257 /// Lazy forward iterator over a map
258 pub struct Entries<'a, K, V> {
259 stack: Vec<&'a TreeNode<K, V>>,
260 // See the comment on MutEntries; this is just to allow
261 // code-sharing (for this immutable-values iterator it *could* very
262 // well be Option<&'a TreeNode<K,V>>).
263 node: *TreeNode<K, V>,
264 remaining_min: uint,
265 remaining_max: uint
266 }
267
268 /// Lazy backward iterator over a map
269 pub struct RevEntries<'a, K, V> {
270 iter: Entries<'a, K, V>,
271 }
272
273 /// Lazy forward iterator over a map that allows for the mutation of
274 /// the values.
275 pub struct MutEntries<'a, K, V> {
276 stack: Vec<&'a mut TreeNode<K, V>>,
277 // Unfortunately, we require some unsafe-ness to get around the
278 // fact that we would be storing a reference *into* one of the
279 // nodes in the stack.
280 //
281 // As far as the compiler knows, this would let us invalidate the
282 // reference by assigning a new value to this node's position in
283 // its parent, which would cause this current one to be
284 // deallocated so this reference would be invalid. (i.e. the
285 // compilers complaints are 100% correct.)
286 //
287 // However, as far as you humans reading this code know (or are
288 // about to know, if you haven't read far enough down yet), we are
289 // only reading from the TreeNode.{left,right} fields. the only
290 // thing that is ever mutated is the .value field (although any
291 // actual mutation that happens is done externally, by the
292 // iterator consumer). So, don't be so concerned, rustc, we've got
293 // it under control.
294 //
295 // (This field can legitimately be null.)
296 node: *mut TreeNode<K, V>,
297 remaining_min: uint,
298 remaining_max: uint
299 }
300
301 /// Lazy backward iterator over a map
302 pub struct RevMutEntries<'a, K, V> {
303 iter: MutEntries<'a, K, V>,
304 }
305
306
307 // FIXME #5846 we want to be able to choose between &x and &mut x
308 // (with many different `x`) below, so we need to optionally pass mut
309 // as a tt, but the only thing we can do with a `tt` is pass them to
310 // other macros, so this takes the `& <mutability> <operand>` token
311 // sequence and forces their evaluation as an expression.
312 macro_rules! addr { ($e:expr) => { $e }}
313 // putting an optional mut into type signatures
314 macro_rules! item { ($i:item) => { $i }}
315
316 macro_rules! define_iterator {
317 ($name:ident,
318 $rev_name:ident,
319
320 // the function to go from &m Option<Box<TreeNode>> to *m TreeNode
321 deref = $deref:ident,
322
323 // see comment on `addr!`, this is just an optional `mut`, but
324 // there's no support for 0-or-1 repeats.
325 addr_mut = $($addr_mut:tt)*
326 ) => {
327 // private methods on the forward iterator (item!() for the
328 // addr_mut in the next_ return value)
329 item!(impl<'a, K, V> $name<'a, K, V> {
330 #[inline(always)]
331 fn next_(&mut self, forward: bool) -> Option<(&'a K, &'a $($addr_mut)* V)> {
332 while !self.stack.is_empty() || !self.node.is_null() {
333 if !self.node.is_null() {
334 let node = unsafe {addr!(& $($addr_mut)* *self.node)};
335 {
336 let next_node = if forward {
337 addr!(& $($addr_mut)* node.left)
338 } else {
339 addr!(& $($addr_mut)* node.right)
340 };
341 self.node = $deref(next_node);
342 }
343 self.stack.push(node);
344 } else {
345 let node = self.stack.pop().unwrap();
346 let next_node = if forward {
347 addr!(& $($addr_mut)* node.right)
348 } else {
349 addr!(& $($addr_mut)* node.left)
350 };
351 self.node = $deref(next_node);
352 self.remaining_max -= 1;
353 if self.remaining_min > 0 {
354 self.remaining_min -= 1;
355 }
356 return Some((&node.key, addr!(& $($addr_mut)* node.value)));
357 }
358 }
359 None
360 }
361
362 /// traverse_left, traverse_right and traverse_complete are
363 /// used to initialize Entries/MutEntries
364 /// pointing to element inside tree structure.
365 ///
366 /// They should be used in following manner:
367 /// - create iterator using TreeMap::[mut_]iter_for_traversal
368 /// - find required node using `traverse_left`/`traverse_right`
369 /// (current node is `Entries::node` field)
370 /// - complete initialization with `traverse_complete`
371 ///
372 /// After this, iteration will start from `self.node`. If
373 /// `self.node` is None iteration will start from last
374 /// node from which we traversed left.
375 #[inline]
376 fn traverse_left(&mut self) {
377 let node = unsafe {addr!(& $($addr_mut)* *self.node)};
378 self.node = $deref(addr!(& $($addr_mut)* node.left));
379 self.stack.push(node);
380 }
381
382 #[inline]
383 fn traverse_right(&mut self) {
384 let node = unsafe {addr!(& $($addr_mut)* *self.node)};
385 self.node = $deref(addr!(& $($addr_mut)* node.right));
386 }
387
388 #[inline]
389 fn traverse_complete(&mut self) {
390 if !self.node.is_null() {
391 unsafe {
392 self.stack.push(addr!(& $($addr_mut)* *self.node));
393 }
394 self.node = ptr::RawPtr::null();
395 }
396 }
397 })
398
399 // the forward Iterator impl.
400 item!(impl<'a, K, V> Iterator<(&'a K, &'a $($addr_mut)* V)> for $name<'a, K, V> {
401 /// Advance the iterator to the next node (in order) and return a
402 /// tuple with a reference to the key and value. If there are no
403 /// more nodes, return `None`.
404 fn next(&mut self) -> Option<(&'a K, &'a $($addr_mut)* V)> {
405 self.next_(true)
406 }
407
408 #[inline]
409 fn size_hint(&self) -> (uint, Option<uint>) {
410 (self.remaining_min, Some(self.remaining_max))
411 }
412 })
413
414 // the reverse Iterator impl.
415 item!(impl<'a, K, V> Iterator<(&'a K, &'a $($addr_mut)* V)> for $rev_name<'a, K, V> {
416 fn next(&mut self) -> Option<(&'a K, &'a $($addr_mut)* V)> {
417 self.iter.next_(false)
418 }
419
420 #[inline]
421 fn size_hint(&self) -> (uint, Option<uint>) {
422 self.iter.size_hint()
423 }
424 })
425 }
426 } // end of define_iterator
427
428 define_iterator! {
429 Entries,
430 RevEntries,
431 deref = deref,
432
433 // immutable, so no mut
434 addr_mut =
435 }
436 define_iterator! {
437 MutEntries,
438 RevMutEntries,
439 deref = mut_deref,
440
441 addr_mut = mut
442 }
443
444 fn deref<'a, K, V>(node: &'a Option<Box<TreeNode<K, V>>>) -> *TreeNode<K, V> {
445 match *node {
446 Some(ref n) => {
447 let n: &TreeNode<K, V> = *n;
448 n as *TreeNode<K, V>
449 }
450 None => ptr::null()
451 }
452 }
453
454 fn mut_deref<K, V>(x: &mut Option<Box<TreeNode<K, V>>>)
455 -> *mut TreeNode<K, V> {
456 match *x {
457 Some(ref mut n) => {
458 let n: &mut TreeNode<K, V> = *n;
459 n as *mut TreeNode<K, V>
460 }
461 None => ptr::mut_null()
462 }
463 }
464
465
466
467 /// Lazy forward iterator over a map that consumes the map while iterating
468 pub struct MoveEntries<K, V> {
469 stack: Vec<TreeNode<K, V>>,
470 remaining: uint
471 }
472
473 impl<K, V> Iterator<(K, V)> for MoveEntries<K,V> {
474 #[inline]
475 fn next(&mut self) -> Option<(K, V)> {
476 while !self.stack.is_empty() {
477 let TreeNode {
478 key: key,
479 value: value,
480 left: left,
481 right: right,
482 level: level
483 } = self.stack.pop().unwrap();
484
485 match left {
486 Some(box left) => {
487 let n = TreeNode {
488 key: key,
489 value: value,
490 left: None,
491 right: right,
492 level: level
493 };
494 self.stack.push(n);
495 self.stack.push(left);
496 }
497 None => {
498 match right {
499 Some(box right) => self.stack.push(right),
500 None => ()
501 }
502 self.remaining -= 1;
503 return Some((key, value))
504 }
505 }
506 }
507 None
508 }
509
510 #[inline]
511 fn size_hint(&self) -> (uint, Option<uint>) {
512 (self.remaining, Some(self.remaining))
513 }
514
515 }
516
517 impl<'a, T> Iterator<&'a T> for SetItems<'a, T> {
518 #[inline]
519 fn next(&mut self) -> Option<&'a T> {
520 self.iter.next().map(|(value, _)| value)
521 }
522 }
523
524 impl<'a, T> Iterator<&'a T> for RevSetItems<'a, T> {
525 #[inline]
526 fn next(&mut self) -> Option<&'a T> {
527 self.iter.next().map(|(value, _)| value)
528 }
529 }
530
531 /// A implementation of the `Set` trait on top of the `TreeMap` container. The
532 /// only requirement is that the type of the elements contained ascribes to the
533 /// `TotalOrd` trait.
534 #[deriving(Clone)]
535 pub struct TreeSet<T> {
536 map: TreeMap<T, ()>
537 }
538
539 impl<T: Eq + TotalOrd> Eq for TreeSet<T> {
540 #[inline]
541 fn eq(&self, other: &TreeSet<T>) -> bool { self.map == other.map }
542 }
543
544 impl<T: Ord + TotalOrd> Ord for TreeSet<T> {
545 #[inline]
546 fn lt(&self, other: &TreeSet<T>) -> bool { self.map < other.map }
547 }
548
549 impl<T: TotalOrd> Container for TreeSet<T> {
550 #[inline]
551 fn len(&self) -> uint { self.map.len() }
552 }
553
554 impl<T: TotalOrd> Mutable for TreeSet<T> {
555 #[inline]
556 fn clear(&mut self) { self.map.clear() }
557 }
558
559 impl<T: TotalOrd> Set<T> for TreeSet<T> {
560 #[inline]
561 fn contains(&self, value: &T) -> bool {
562 self.map.contains_key(value)
563 }
564
565 fn is_disjoint(&self, other: &TreeSet<T>) -> bool {
566 self.intersection(other).next().is_none()
567 }
568
569 fn is_subset(&self, other: &TreeSet<T>) -> bool {
570 let mut x = self.iter();
571 let mut y = other.iter();
572 let mut a = x.next();
573 let mut b = y.next();
574 while a.is_some() {
575 if b.is_none() {
576 return false;
577 }
578
579 let a1 = a.unwrap();
580 let b1 = b.unwrap();
581
582 match b1.cmp(a1) {
583 Less => (),
584 Greater => return false,
585 Equal => a = x.next(),
586 }
587
588 b = y.next();
589 }
590 true
591 }
592 }
593
594 impl<T: TotalOrd> MutableSet<T> for TreeSet<T> {
595 #[inline]
596 fn insert(&mut self, value: T) -> bool { self.map.insert(value, ()) }
597
598 #[inline]
599 fn remove(&mut self, value: &T) -> bool { self.map.remove(value) }
600 }
601
602 impl<T: TotalOrd> TreeSet<T> {
603 /// Create an empty TreeSet
604 #[inline]
605 pub fn new() -> TreeSet<T> { TreeSet{map: TreeMap::new()} }
606
607 /// Get a lazy iterator over the values in the set.
608 /// Requires that it be frozen (immutable).
609 #[inline]
610 pub fn iter<'a>(&'a self) -> SetItems<'a, T> {
611 SetItems{iter: self.map.iter()}
612 }
613
614 /// Get a lazy iterator over the values in the set.
615 /// Requires that it be frozen (immutable).
616 #[inline]
617 pub fn rev_iter<'a>(&'a self) -> RevSetItems<'a, T> {
618 RevSetItems{iter: self.map.rev_iter()}
619 }
620
621 /// Get a lazy iterator pointing to the first value not less than `v` (greater or equal).
622 /// If all elements in the set are less than `v` empty iterator is returned.
623 #[inline]
624 pub fn lower_bound<'a>(&'a self, v: &T) -> SetItems<'a, T> {
625 SetItems{iter: self.map.lower_bound(v)}
626 }
627
628 /// Get a lazy iterator pointing to the first value greater than `v`.
629 /// If all elements in the set are not greater than `v` empty iterator is returned.
630 #[inline]
631 pub fn upper_bound<'a>(&'a self, v: &T) -> SetItems<'a, T> {
632 SetItems{iter: self.map.upper_bound(v)}
633 }
634
635 /// Visit the values (in-order) representing the difference
636 pub fn difference<'a>(&'a self, other: &'a TreeSet<T>) -> DifferenceItems<'a, T> {
637 DifferenceItems{a: self.iter().peekable(), b: other.iter().peekable()}
638 }
639
640 /// Visit the values (in-order) representing the symmetric difference
641 pub fn symmetric_difference<'a>(&'a self, other: &'a TreeSet<T>)
642 -> SymDifferenceItems<'a, T> {
643 SymDifferenceItems{a: self.iter().peekable(), b: other.iter().peekable()}
644 }
645
646 /// Visit the values (in-order) representing the intersection
647 pub fn intersection<'a>(&'a self, other: &'a TreeSet<T>)
648 -> IntersectionItems<'a, T> {
649 IntersectionItems{a: self.iter().peekable(), b: other.iter().peekable()}
650 }
651
652 /// Visit the values (in-order) representing the union
653 pub fn union<'a>(&'a self, other: &'a TreeSet<T>) -> UnionItems<'a, T> {
654 UnionItems{a: self.iter().peekable(), b: other.iter().peekable()}
655 }
656 }
657
658 /// Lazy forward iterator over a set
659 pub struct SetItems<'a, T> {
660 iter: Entries<'a, T, ()>
661 }
662
663 /// Lazy backward iterator over a set
664 pub struct RevSetItems<'a, T> {
665 iter: RevEntries<'a, T, ()>
666 }
667
668 /// Lazy iterator producing elements in the set difference (in-order)
669 pub struct DifferenceItems<'a, T> {
670 a: Peekable<&'a T, SetItems<'a, T>>,
671 b: Peekable<&'a T, SetItems<'a, T>>,
672 }
673
674 /// Lazy iterator producing elements in the set symmetric difference (in-order)
675 pub struct SymDifferenceItems<'a, T> {
676 a: Peekable<&'a T, SetItems<'a, T>>,
677 b: Peekable<&'a T, SetItems<'a, T>>,
678 }
679
680 /// Lazy iterator producing elements in the set intersection (in-order)
681 pub struct IntersectionItems<'a, T> {
682 a: Peekable<&'a T, SetItems<'a, T>>,
683 b: Peekable<&'a T, SetItems<'a, T>>,
684 }
685
686 /// Lazy iterator producing elements in the set intersection (in-order)
687 pub struct UnionItems<'a, T> {
688 a: Peekable<&'a T, SetItems<'a, T>>,
689 b: Peekable<&'a T, SetItems<'a, T>>,
690 }
691
692 /// Compare `x` and `y`, but return `short` if x is None and `long` if y is None
693 fn cmp_opt<T: TotalOrd>(x: Option<&T>, y: Option<&T>,
694 short: Ordering, long: Ordering) -> Ordering {
695 match (x, y) {
696 (None , _ ) => short,
697 (_ , None ) => long,
698 (Some(x1), Some(y1)) => x1.cmp(y1),
699 }
700 }
701
702 impl<'a, T: TotalOrd> Iterator<&'a T> for DifferenceItems<'a, T> {
703 fn next(&mut self) -> Option<&'a T> {
704 loop {
705 match cmp_opt(self.a.peek(), self.b.peek(), Less, Less) {
706 Less => return self.a.next(),
707 Equal => { self.a.next(); self.b.next(); }
708 Greater => { self.b.next(); }
709 }
710 }
711 }
712 }
713
714 impl<'a, T: TotalOrd> Iterator<&'a T> for SymDifferenceItems<'a, T> {
715 fn next(&mut self) -> Option<&'a T> {
716 loop {
717 match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) {
718 Less => return self.a.next(),
719 Equal => { self.a.next(); self.b.next(); }
720 Greater => return self.b.next(),
721 }
722 }
723 }
724 }
725
726 impl<'a, T: TotalOrd> Iterator<&'a T> for IntersectionItems<'a, T> {
727 fn next(&mut self) -> Option<&'a T> {
728 loop {
729 let o_cmp = match (self.a.peek(), self.b.peek()) {
730 (None , _ ) => None,
731 (_ , None ) => None,
732 (Some(a1), Some(b1)) => Some(a1.cmp(b1)),
733 };
734 match o_cmp {
735 None => return None,
736 Some(Less) => { self.a.next(); }
737 Some(Equal) => { self.b.next(); return self.a.next() }
738 Some(Greater) => { self.b.next(); }
739 }
740 }
741 }
742 }
743
744 impl<'a, T: TotalOrd> Iterator<&'a T> for UnionItems<'a, T> {
745 fn next(&mut self) -> Option<&'a T> {
746 loop {
747 match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) {
748 Less => return self.a.next(),
749 Equal => { self.b.next(); return self.a.next() }
750 Greater => return self.b.next(),
751 }
752 }
753 }
754 }
755
756
757 // Nodes keep track of their level in the tree, starting at 1 in the
758 // leaves and with a red child sharing the level of the parent.
759 #[deriving(Clone)]
760 struct TreeNode<K, V> {
761 key: K,
762 value: V,
763 left: Option<Box<TreeNode<K, V>>>,
764 right: Option<Box<TreeNode<K, V>>>,
765 level: uint
766 }
767
768 impl<K: TotalOrd, V> TreeNode<K, V> {
769 /// Creates a new tree node.
770 #[inline]
771 pub fn new(key: K, value: V) -> TreeNode<K, V> {
772 TreeNode{key: key, value: value, left: None, right: None, level: 1}
773 }
774 }
775
776 // Remove left horizontal link by rotating right
777 fn skew<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>) {
778 if node.left.as_ref().map_or(false, |x| x.level == node.level) {
779 let mut save = node.left.take_unwrap();
780 swap(&mut node.left, &mut save.right); // save.right now None
781 swap(node, &mut save);
782 node.right = Some(save);
783 }
784 }
785
786 // Remove dual horizontal link by rotating left and increasing level of
787 // the parent
788 fn split<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>) {
789 if node.right.as_ref().map_or(false,
790 |x| x.right.as_ref().map_or(false, |y| y.level == node.level)) {
791 let mut save = node.right.take_unwrap();
792 swap(&mut node.right, &mut save.left); // save.left now None
793 save.level += 1;
794 swap(node, &mut save);
795 node.left = Some(save);
796 }
797 }
798
799 fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<Box<TreeNode<K, V>>>,
800 key: &K)
801 -> Option<&'r mut V> {
802 match *node {
803 Some(ref mut x) => {
804 match key.cmp(&x.key) {
805 Less => find_mut(&mut x.left, key),
806 Greater => find_mut(&mut x.right, key),
807 Equal => Some(&mut x.value),
808 }
809 }
810 None => None
811 }
812 }
813
814 fn insert<K: TotalOrd, V>(node: &mut Option<Box<TreeNode<K, V>>>,
815 key: K, value: V) -> Option<V> {
816 match *node {
817 Some(ref mut save) => {
818 match key.cmp(&save.key) {
819 Less => {
820 let inserted = insert(&mut save.left, key, value);
821 skew(save);
822 split(save);
823 inserted
824 }
825 Greater => {
826 let inserted = insert(&mut save.right, key, value);
827 skew(save);
828 split(save);
829 inserted
830 }
831 Equal => {
832 save.key = key;
833 Some(replace(&mut save.value, value))
834 }
835 }
836 }
837 None => {
838 *node = Some(box TreeNode::new(key, value));
839 None
840 }
841 }
842 }
843
844 fn remove<K: TotalOrd, V>(node: &mut Option<Box<TreeNode<K, V>>>,
845 key: &K) -> Option<V> {
846 fn heir_swap<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>,
847 child: &mut Option<Box<TreeNode<K, V>>>) {
848 // *could* be done without recursion, but it won't borrow check
849 for x in child.mut_iter() {
850 if x.right.is_some() {
851 heir_swap(node, &mut x.right);
852 } else {
853 swap(&mut node.key, &mut x.key);
854 swap(&mut node.value, &mut x.value);
855 }
856 }
857 }
858
859 match *node {
860 None => {
861 return None; // bottom of tree
862 }
863 Some(ref mut save) => {
864 let (ret, rebalance) = match key.cmp(&save.key) {
865 Less => (remove(&mut save.left, key), true),
866 Greater => (remove(&mut save.right, key), true),
867 Equal => {
868 if save.left.is_some() {
869 if save.right.is_some() {
870 let mut left = save.left.take_unwrap();
871 if left.right.is_some() {
872 heir_swap(save, &mut left.right);
873 } else {
874 swap(&mut save.key, &mut left.key);
875 swap(&mut save.value, &mut left.value);
876 }
877 save.left = Some(left);
878 (remove(&mut save.left, key), true)
879 } else {
880 let new = save.left.take_unwrap();
881 let box TreeNode{value, ..} = replace(save, new);
882 *save = save.left.take_unwrap();
883 (Some(value), true)
884 }
885 } else if save.right.is_some() {
886 let new = save.right.take_unwrap();
887 let box TreeNode{value, ..} = replace(save, new);
888 (Some(value), true)
889 } else {
890 (None, false)
891 }
892 }
893 };
894
895 if rebalance {
896 let left_level = save.left.as_ref().map_or(0, |x| x.level);
897 let right_level = save.right.as_ref().map_or(0, |x| x.level);
898
899 // re-balance, if necessary
900 if left_level < save.level - 1 || right_level < save.level - 1 {
901 save.level -= 1;
902
903 if right_level > save.level {
904 for x in save.right.mut_iter() { x.level = save.level }
905 }
906
907 skew(save);
908
909 for right in save.right.mut_iter() {
910 skew(right);
911 for x in right.right.mut_iter() { skew(x) }
912 }
913
914 split(save);
915 for x in save.right.mut_iter() { split(x) }
916 }
917
918 return ret;
919 }
920 }
921 }
922 return match node.take() {
923 Some(box TreeNode{value, ..}) => Some(value), None => fail!()
924 };
925 }
926
927 impl<K: TotalOrd, V> FromIterator<(K, V)> for TreeMap<K, V> {
928 fn from_iter<T: Iterator<(K, V)>>(iter: T) -> TreeMap<K, V> {
929 let mut map = TreeMap::new();
930 map.extend(iter);
931 map
932 }
933 }
934
935 impl<K: TotalOrd, V> Extendable<(K, V)> for TreeMap<K, V> {
936 #[inline]
937 fn extend<T: Iterator<(K, V)>>(&mut self, mut iter: T) {
938 for (k, v) in iter {
939 self.insert(k, v);
940 }
941 }
942 }
943
944 impl<T: TotalOrd> FromIterator<T> for TreeSet<T> {
945 fn from_iter<Iter: Iterator<T>>(iter: Iter) -> TreeSet<T> {
946 let mut set = TreeSet::new();
947 set.extend(iter);
948 set
949 }
950 }
951
952 impl<T: TotalOrd> Extendable<T> for TreeSet<T> {
953 #[inline]
954 fn extend<Iter: Iterator<T>>(&mut self, mut iter: Iter) {
955 for elem in iter {
956 self.insert(elem);
957 }
958 }
959 }
960
961 #[cfg(test)]
962 mod test_treemap {
963 use super::{TreeMap, TreeNode};
964
965 use rand::Rng;
966 use rand;
967
968 #[test]
969 fn find_empty() {
970 let m: TreeMap<int,int> = TreeMap::new();
971 assert!(m.find(&5) == None);
972 }
973
974 #[test]
975 fn find_not_found() {
976 let mut m = TreeMap::new();
977 assert!(m.insert(1, 2));
978 assert!(m.insert(5, 3));
979 assert!(m.insert(9, 3));
980 assert_eq!(m.find(&2), None);
981 }
982
983 #[test]
984 fn test_find_mut() {
985 let mut m = TreeMap::new();
986 assert!(m.insert(1, 12));
987 assert!(m.insert(2, 8));
988 assert!(m.insert(5, 14));
989 let new = 100;
990 match m.find_mut(&5) {
991 None => fail!(), Some(x) => *x = new
992 }
993 assert_eq!(m.find(&5), Some(&new));
994 }
995
996 #[test]
997 fn insert_replace() {
998 let mut m = TreeMap::new();
999 assert!(m.insert(5, 2));
1000 assert!(m.insert(2, 9));
1001 assert!(!m.insert(2, 11));
1002 assert_eq!(m.find(&2).unwrap(), &11);
1003 }
1004
1005 #[test]
1006 fn test_clear() {
1007 let mut m = TreeMap::new();
1008 m.clear();
1009 assert!(m.insert(5, 11));
1010 assert!(m.insert(12, -3));
1011 assert!(m.insert(19, 2));
1012 m.clear();
1013 assert!(m.find(&5).is_none());
1014 assert!(m.find(&12).is_none());
1015 assert!(m.find(&19).is_none());
1016 assert!(m.is_empty());
1017 }
1018
1019 #[test]
1020 fn u8_map() {
1021 let mut m = TreeMap::new();
1022
1023 let k1 = "foo".as_bytes();
1024 let k2 = "bar".as_bytes();
1025 let v1 = "baz".as_bytes();
1026 let v2 = "foobar".as_bytes();
1027
1028 m.insert(k1.clone(), v1.clone());
1029 m.insert(k2.clone(), v2.clone());
1030
1031 assert_eq!(m.find(&k2), Some(&v2));
1032 assert_eq!(m.find(&k1), Some(&v1));
1033 }
1034
1035 fn check_equal<K: Eq + TotalOrd, V: Eq>(ctrl: &[(K, V)],
1036 map: &TreeMap<K, V>) {
1037 assert_eq!(ctrl.is_empty(), map.is_empty());
1038 for x in ctrl.iter() {
1039 let &(ref k, ref v) = x;
1040 assert!(map.find(k).unwrap() == v)
1041 }
1042 for (map_k, map_v) in map.iter() {
1043 let mut found = false;
1044 for x in ctrl.iter() {
1045 let &(ref ctrl_k, ref ctrl_v) = x;
1046 if *map_k == *ctrl_k {
1047 assert!(*map_v == *ctrl_v);
1048 found = true;
1049 break;
1050 }
1051 }
1052 assert!(found);
1053 }
1054 }
1055
1056 fn check_left<K: TotalOrd, V>(node: &Option<Box<TreeNode<K, V>>>,
1057 parent: &Box<TreeNode<K, V>>) {
1058 match *node {
1059 Some(ref r) => {
1060 assert_eq!(r.key.cmp(&parent.key), Less);
1061 assert!(r.level == parent.level - 1); // left is black
1062 check_left(&r.left, r);
1063 check_right(&r.right, r, false);
1064 }
1065 None => assert!(parent.level == 1) // parent is leaf
1066 }
1067 }
1068
1069 fn check_right<K: TotalOrd, V>(node: &Option<Box<TreeNode<K, V>>>,
1070 parent: &Box<TreeNode<K, V>>,
1071 parent_red: bool) {
1072 match *node {
1073 Some(ref r) => {
1074 assert_eq!(r.key.cmp(&parent.key), Greater);
1075 let red = r.level == parent.level;
1076 if parent_red { assert!(!red) } // no dual horizontal links
1077 // Right red or black
1078 assert!(red || r.level == parent.level - 1);
1079 check_left(&r.left, r);
1080 check_right(&r.right, r, red);
1081 }
1082 None => assert!(parent.level == 1) // parent is leaf
1083 }
1084 }
1085
1086 fn check_structure<K: TotalOrd, V>(map: &TreeMap<K, V>) {
1087 match map.root {
1088 Some(ref r) => {
1089 check_left(&r.left, r);
1090 check_right(&r.right, r, false);
1091 }
1092 None => ()
1093 }
1094 }
1095
1096 #[test]
1097 fn test_rand_int() {
1098 let mut map: TreeMap<int,int> = TreeMap::new();
1099 let mut ctrl = vec![];
1100
1101 check_equal(ctrl.as_slice(), &map);
1102 assert!(map.find(&5).is_none());
1103
1104 let mut rng: rand::IsaacRng = rand::SeedableRng::from_seed(&[42]);
1105
1106 for _ in range(0, 3) {
1107 for _ in range(0, 90) {
1108 let k = rng.gen();
1109 let v = rng.gen();
1110 if !ctrl.iter().any(|x| x == &(k, v)) {
1111 assert!(map.insert(k, v));
1112 ctrl.push((k, v));
1113 check_structure(&map);
1114 check_equal(ctrl.as_slice(), &map);
1115 }
1116 }
1117
1118 for _ in range(0, 30) {
1119 let r = rng.gen_range(0, ctrl.len());
1120 let (key, _) = ctrl.remove(r).unwrap();
1121 assert!(map.remove(&key));
1122 check_structure(&map);
1123 check_equal(ctrl.as_slice(), &map);
1124 }
1125 }
1126 }
1127
1128 #[test]
1129 fn test_len() {
1130 let mut m = TreeMap::new();
1131 assert!(m.insert(3, 6));
1132 assert_eq!(m.len(), 1);
1133 assert!(m.insert(0, 0));
1134 assert_eq!(m.len(), 2);
1135 assert!(m.insert(4, 8));
1136 assert_eq!(m.len(), 3);
1137 assert!(m.remove(&3));
1138 assert_eq!(m.len(), 2);
1139 assert!(!m.remove(&5));
1140 assert_eq!(m.len(), 2);
1141 assert!(m.insert(2, 4));
1142 assert_eq!(m.len(), 3);
1143 assert!(m.insert(1, 2));
1144 assert_eq!(m.len(), 4);
1145 }
1146
1147 #[test]
1148 fn test_iterator() {
1149 let mut m = TreeMap::new();
1150
1151 assert!(m.insert(3, 6));
1152 assert!(m.insert(0, 0));
1153 assert!(m.insert(4, 8));
1154 assert!(m.insert(2, 4));
1155 assert!(m.insert(1, 2));
1156
1157 let mut n = 0;
1158 for (k, v) in m.iter() {
1159 assert_eq!(*k, n);
1160 assert_eq!(*v, n * 2);
1161 n += 1;
1162 }
1163 assert_eq!(n, 5);
1164 }
1165
1166 #[test]
1167 fn test_interval_iteration() {
1168 let mut m = TreeMap::new();
1169 for i in range(1, 100) {
1170 assert!(m.insert(i * 2, i * 4));
1171 }
1172
1173 for i in range(1, 198) {
1174 let mut lb_it = m.lower_bound(&i);
1175 let (&k, &v) = lb_it.next().unwrap();
1176 let lb = i + i % 2;
1177 assert_eq!(lb, k);
1178 assert_eq!(lb * 2, v);
1179
1180 let mut ub_it = m.upper_bound(&i);
1181 let (&k, &v) = ub_it.next().unwrap();
1182 let ub = i + 2 - i % 2;
1183 assert_eq!(ub, k);
1184 assert_eq!(ub * 2, v);
1185 }
1186 let mut end_it = m.lower_bound(&199);
1187 assert_eq!(end_it.next(), None);
1188 }
1189
1190 #[test]
1191 fn test_rev_iter() {
1192 let mut m = TreeMap::new();
1193
1194 assert!(m.insert(3, 6));
1195 assert!(m.insert(0, 0));
1196 assert!(m.insert(4, 8));
1197 assert!(m.insert(2, 4));
1198 assert!(m.insert(1, 2));
1199
1200 let mut n = 4;
1201 for (k, v) in m.rev_iter() {
1202 assert_eq!(*k, n);
1203 assert_eq!(*v, n * 2);
1204 n -= 1;
1205 }
1206 }
1207
1208 #[test]
1209 fn test_mut_iter() {
1210 let mut m = TreeMap::new();
1211 for i in range(0u, 10) {
1212 assert!(m.insert(i, 100 * i));
1213 }
1214
1215 for (i, (&k, v)) in m.mut_iter().enumerate() {
1216 *v += k * 10 + i; // 000 + 00 + 0, 100 + 10 + 1, ...
1217 }
1218
1219 for (&k, &v) in m.iter() {
1220 assert_eq!(v, 111 * k);
1221 }
1222 }
1223 #[test]
1224 fn test_mut_rev_iter() {
1225 let mut m = TreeMap::new();
1226 for i in range(0u, 10) {
1227 assert!(m.insert(i, 100 * i));
1228 }
1229
1230 for (i, (&k, v)) in m.mut_rev_iter().enumerate() {
1231 *v += k * 10 + (9 - i); // 900 + 90 + (9 - 0), 800 + 80 + (9 - 1), ...
1232 }
1233
1234 for (&k, &v) in m.iter() {
1235 assert_eq!(v, 111 * k);
1236 }
1237 }
1238
1239 #[test]
1240 fn test_mut_interval_iter() {
1241 let mut m_lower = TreeMap::new();
1242 let mut m_upper = TreeMap::new();
1243 for i in range(1, 100) {
1244 assert!(m_lower.insert(i * 2, i * 4));
1245 assert!(m_upper.insert(i * 2, i * 4));
1246 }
1247
1248 for i in range(1, 199) {
1249 let mut lb_it = m_lower.mut_lower_bound(&i);
1250 let (&k, v) = lb_it.next().unwrap();
1251 let lb = i + i % 2;
1252 assert_eq!(lb, k);
1253 *v -= k;
1254 }
1255 for i in range(0, 198) {
1256 let mut ub_it = m_upper.mut_upper_bound(&i);
1257 let (&k, v) = ub_it.next().unwrap();
1258 let ub = i + 2 - i % 2;
1259 assert_eq!(ub, k);
1260 *v -= k;
1261 }
1262
1263 assert!(m_lower.mut_lower_bound(&199).next().is_none());
1264
1265 assert!(m_upper.mut_upper_bound(&198).next().is_none());
1266
1267 assert!(m_lower.iter().all(|(_, &x)| x == 0));
1268 assert!(m_upper.iter().all(|(_, &x)| x == 0));
1269 }
1270
1271 #[test]
1272 fn test_eq() {
1273 let mut a = TreeMap::new();
1274 let mut b = TreeMap::new();
1275
1276 assert!(a == b);
1277 assert!(a.insert(0, 5));
1278 assert!(a != b);
1279 assert!(b.insert(0, 4));
1280 assert!(a != b);
1281 assert!(a.insert(5, 19));
1282 assert!(a != b);
1283 assert!(!b.insert(0, 5));
1284 assert!(a != b);
1285 assert!(b.insert(5, 19));
1286 assert!(a == b);
1287 }
1288
1289 #[test]
1290 fn test_lt() {
1291 let mut a = TreeMap::new();
1292 let mut b = TreeMap::new();
1293
1294 assert!(!(a < b) && !(b < a));
1295 assert!(b.insert(0, 5));
1296 assert!(a < b);
1297 assert!(a.insert(0, 7));
1298 assert!(!(a < b) && b < a);
1299 assert!(b.insert(-2, 0));
1300 assert!(b < a);
1301 assert!(a.insert(-5, 2));
1302 assert!(a < b);
1303 assert!(a.insert(6, 2));
1304 assert!(a < b && !(b < a));
1305 }
1306
1307 #[test]
1308 fn test_ord() {
1309 let mut a = TreeMap::new();
1310 let mut b = TreeMap::new();
1311
1312 assert!(a <= b && a >= b);
1313 assert!(a.insert(1, 1));
1314 assert!(a > b && a >= b);
1315 assert!(b < a && b <= a);
1316 assert!(b.insert(2, 2));
1317 assert!(b > a && b >= a);
1318 assert!(a < b && a <= b);
1319 }
1320
1321 #[test]
1322 fn test_lazy_iterator() {
1323 let mut m = TreeMap::new();
1324 let (x1, y1) = (2, 5);
1325 let (x2, y2) = (9, 12);
1326 let (x3, y3) = (20, -3);
1327 let (x4, y4) = (29, 5);
1328 let (x5, y5) = (103, 3);
1329
1330 assert!(m.insert(x1, y1));
1331 assert!(m.insert(x2, y2));
1332 assert!(m.insert(x3, y3));
1333 assert!(m.insert(x4, y4));
1334 assert!(m.insert(x5, y5));
1335
1336 let m = m;
1337 let mut a = m.iter();
1338
1339 assert_eq!(a.next().unwrap(), (&x1, &y1));
1340 assert_eq!(a.next().unwrap(), (&x2, &y2));
1341 assert_eq!(a.next().unwrap(), (&x3, &y3));
1342 assert_eq!(a.next().unwrap(), (&x4, &y4));
1343 assert_eq!(a.next().unwrap(), (&x5, &y5));
1344
1345 assert!(a.next().is_none());
1346
1347 let mut b = m.iter();
1348
1349 let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4),
1350 (&x5, &y5)];
1351 let mut i = 0;
1352
1353 for x in b {
1354 assert_eq!(expected[i], x);
1355 i += 1;
1356
1357 if i == 2 {
1358 break
1359 }
1360 }
1361
1362 for x in b {
1363 assert_eq!(expected[i], x);
1364 i += 1;
1365 }
1366 }
1367
1368 #[test]
1369 fn test_from_iter() {
1370 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1371
1372 let map: TreeMap<int, int> = xs.iter().map(|&x| x).collect();
1373
1374 for &(k, v) in xs.iter() {
1375 assert_eq!(map.find(&k), Some(&v));
1376 }
1377 }
1378
1379 }
1380
1381 #[cfg(test)]
1382 mod bench {
1383 extern crate test;
1384 use self::test::Bencher;
1385 use super::TreeMap;
1386 use deque::bench::{insert_rand_n, insert_seq_n, find_rand_n, find_seq_n};
1387
1388 // Find seq
1389 #[bench]
1390 pub fn insert_rand_100(b: &mut Bencher) {
1391 let mut m : TreeMap<uint,uint> = TreeMap::new();
1392 insert_rand_n(100, &mut m, b);
1393 }
1394
1395 #[bench]
1396 pub fn insert_rand_10_000(b: &mut Bencher) {
1397 let mut m : TreeMap<uint,uint> = TreeMap::new();
1398 insert_rand_n(10_000, &mut m, b);
1399 }
1400
1401 // Insert seq
1402 #[bench]
1403 pub fn insert_seq_100(b: &mut Bencher) {
1404 let mut m : TreeMap<uint,uint> = TreeMap::new();
1405 insert_seq_n(100, &mut m, b);
1406 }
1407
1408 #[bench]
1409 pub fn insert_seq_10_000(b: &mut Bencher) {
1410 let mut m : TreeMap<uint,uint> = TreeMap::new();
1411 insert_seq_n(10_000, &mut m, b);
1412 }
1413
1414 // Find rand
1415 #[bench]
1416 pub fn find_rand_100(b: &mut Bencher) {
1417 let mut m : TreeMap<uint,uint> = TreeMap::new();
1418 find_rand_n(100, &mut m, b);
1419 }
1420
1421 #[bench]
1422 pub fn find_rand_10_000(b: &mut Bencher) {
1423 let mut m : TreeMap<uint,uint> = TreeMap::new();
1424 find_rand_n(10_000, &mut m, b);
1425 }
1426
1427 // Find seq
1428 #[bench]
1429 pub fn find_seq_100(b: &mut Bencher) {
1430 let mut m : TreeMap<uint,uint> = TreeMap::new();
1431 find_seq_n(100, &mut m, b);
1432 }
1433
1434 #[bench]
1435 pub fn find_seq_10_000(b: &mut Bencher) {
1436 let mut m : TreeMap<uint,uint> = TreeMap::new();
1437 find_seq_n(10_000, &mut m, b);
1438 }
1439 }
1440
1441 #[cfg(test)]
1442 mod test_set {
1443
1444 use super::{TreeMap, TreeSet};
1445
1446 #[test]
1447 fn test_clear() {
1448 let mut s = TreeSet::new();
1449 s.clear();
1450 assert!(s.insert(5));
1451 assert!(s.insert(12));
1452 assert!(s.insert(19));
1453 s.clear();
1454 assert!(!s.contains(&5));
1455 assert!(!s.contains(&12));
1456 assert!(!s.contains(&19));
1457 assert!(s.is_empty());
1458 }
1459
1460 #[test]
1461 fn test_disjoint() {
1462 let mut xs = TreeSet::new();
1463 let mut ys = TreeSet::new();
1464 assert!(xs.is_disjoint(&ys));
1465 assert!(ys.is_disjoint(&xs));
1466 assert!(xs.insert(5));
1467 assert!(ys.insert(11));
1468 assert!(xs.is_disjoint(&ys));
1469 assert!(ys.is_disjoint(&xs));
1470 assert!(xs.insert(7));
1471 assert!(xs.insert(19));
1472 assert!(xs.insert(4));
1473 assert!(ys.insert(2));
1474 assert!(ys.insert(-11));
1475 assert!(xs.is_disjoint(&ys));
1476 assert!(ys.is_disjoint(&xs));
1477 assert!(ys.insert(7));
1478 assert!(!xs.is_disjoint(&ys));
1479 assert!(!ys.is_disjoint(&xs));
1480 }
1481
1482 #[test]
1483 fn test_subset_and_superset() {
1484 let mut a = TreeSet::new();
1485 assert!(a.insert(0));
1486 assert!(a.insert(5));
1487 assert!(a.insert(11));
1488 assert!(a.insert(7));
1489
1490 let mut b = TreeSet::new();
1491 assert!(b.insert(0));
1492 assert!(b.insert(7));
1493 assert!(b.insert(19));
1494 assert!(b.insert(250));
1495 assert!(b.insert(11));
1496 assert!(b.insert(200));
1497
1498 assert!(!a.is_subset(&b));
1499 assert!(!a.is_superset(&b));
1500 assert!(!b.is_subset(&a));
1501 assert!(!b.is_superset(&a));
1502
1503 assert!(b.insert(5));
1504
1505 assert!(a.is_subset(&b));
1506 assert!(!a.is_superset(&b));
1507 assert!(!b.is_subset(&a));
1508 assert!(b.is_superset(&a));
1509 }
1510
1511 #[test]
1512 fn test_iterator() {
1513 let mut m = TreeSet::new();
1514
1515 assert!(m.insert(3));
1516 assert!(m.insert(0));
1517 assert!(m.insert(4));
1518 assert!(m.insert(2));
1519 assert!(m.insert(1));
1520
1521 let mut n = 0;
1522 for x in m.iter() {
1523 assert_eq!(*x, n);
1524 n += 1
1525 }
1526 }
1527
1528 #[test]
1529 fn test_rev_iter() {
1530 let mut m = TreeSet::new();
1531
1532 assert!(m.insert(3));
1533 assert!(m.insert(0));
1534 assert!(m.insert(4));
1535 assert!(m.insert(2));
1536 assert!(m.insert(1));
1537
1538 let mut n = 4;
1539 for x in m.rev_iter() {
1540 assert_eq!(*x, n);
1541 n -= 1;
1542 }
1543 }
1544
1545 #[test]
1546 fn test_clone_eq() {
1547 let mut m = TreeSet::new();
1548
1549 m.insert(1);
1550 m.insert(2);
1551
1552 assert!(m.clone() == m);
1553 }
1554
1555 fn check(a: &[int],
1556 b: &[int],
1557 expected: &[int],
1558 f: |&TreeSet<int>, &TreeSet<int>, f: |&int| -> bool| -> bool) {
1559 let mut set_a = TreeSet::new();
1560 let mut set_b = TreeSet::new();
1561
1562 for x in a.iter() { assert!(set_a.insert(*x)) }
1563 for y in b.iter() { assert!(set_b.insert(*y)) }
1564
1565 let mut i = 0;
1566 f(&set_a, &set_b, |x| {
1567 assert_eq!(*x, expected[i]);
1568 i += 1;
1569 true
1570 });
1571 assert_eq!(i, expected.len());
1572 }
1573
1574 #[test]
1575 fn test_intersection() {
1576 fn check_intersection(a: &[int], b: &[int], expected: &[int]) {
1577 check(a, b, expected, |x, y, f| x.intersection(y).advance(f))
1578 }
1579
1580 check_intersection([], [], []);
1581 check_intersection([1, 2, 3], [], []);
1582 check_intersection([], [1, 2, 3], []);
1583 check_intersection([2], [1, 2, 3], [2]);
1584 check_intersection([1, 2, 3], [2], [2]);
1585 check_intersection([11, 1, 3, 77, 103, 5, -5],
1586 [2, 11, 77, -9, -42, 5, 3],
1587 [3, 5, 11, 77]);
1588 }
1589
1590 #[test]
1591 fn test_difference() {
1592 fn check_difference(a: &[int], b: &[int], expected: &[int]) {
1593 check(a, b, expected, |x, y, f| x.difference(y).advance(f))
1594 }
1595
1596 check_difference([], [], []);
1597 check_difference([1, 12], [], [1, 12]);
1598 check_difference([], [1, 2, 3, 9], []);
1599 check_difference([1, 3, 5, 9, 11],
1600 [3, 9],
1601 [1, 5, 11]);
1602 check_difference([-5, 11, 22, 33, 40, 42],
1603 [-12, -5, 14, 23, 34, 38, 39, 50],
1604 [11, 22, 33, 40, 42]);
1605 }
1606
1607 #[test]
1608 fn test_symmetric_difference() {
1609 fn check_symmetric_difference(a: &[int], b: &[int],
1610 expected: &[int]) {
1611 check(a, b, expected, |x, y, f| x.symmetric_difference(y).advance(f))
1612 }
1613
1614 check_symmetric_difference([], [], []);
1615 check_symmetric_difference([1, 2, 3], [2], [1, 3]);
1616 check_symmetric_difference([2], [1, 2, 3], [1, 3]);
1617 check_symmetric_difference([1, 3, 5, 9, 11],
1618 [-2, 3, 9, 14, 22],
1619 [-2, 1, 5, 11, 14, 22]);
1620 }
1621
1622 #[test]
1623 fn test_union() {
1624 fn check_union(a: &[int], b: &[int],
1625 expected: &[int]) {
1626 check(a, b, expected, |x, y, f| x.union(y).advance(f))
1627 }
1628
1629 check_union([], [], []);
1630 check_union([1, 2, 3], [2], [1, 2, 3]);
1631 check_union([2], [1, 2, 3], [1, 2, 3]);
1632 check_union([1, 3, 5, 9, 11, 16, 19, 24],
1633 [-2, 1, 5, 9, 13, 19],
1634 [-2, 1, 3, 5, 9, 11, 13, 16, 19, 24]);
1635 }
1636
1637 #[test]
1638 fn test_zip() {
1639 let mut x = TreeSet::new();
1640 x.insert(5u);
1641 x.insert(12u);
1642 x.insert(11u);
1643
1644 let mut y = TreeSet::new();
1645 y.insert("foo");
1646 y.insert("bar");
1647
1648 let x = x;
1649 let y = y;
1650 let mut z = x.iter().zip(y.iter());
1651
1652 // FIXME: #5801: this needs a type hint to compile...
1653 let result: Option<(&uint, & &'static str)> = z.next();
1654 assert_eq!(result.unwrap(), (&5u, &("bar")));
1655
1656 let result: Option<(&uint, & &'static str)> = z.next();
1657 assert_eq!(result.unwrap(), (&11u, &("foo")));
1658
1659 let result: Option<(&uint, & &'static str)> = z.next();
1660 assert!(result.is_none());
1661 }
1662
1663 #[test]
1664 fn test_swap() {
1665 let mut m = TreeMap::new();
1666 assert_eq!(m.swap(1, 2), None);
1667 assert_eq!(m.swap(1, 3), Some(2));
1668 assert_eq!(m.swap(1, 4), Some(3));
1669 }
1670
1671 #[test]
1672 fn test_pop() {
1673 let mut m = TreeMap::new();
1674 m.insert(1, 2);
1675 assert_eq!(m.pop(&1), Some(2));
1676 assert_eq!(m.pop(&1), None);
1677 }
1678
1679 #[test]
1680 fn test_from_iter() {
1681 let xs = [1, 2, 3, 4, 5, 6, 7, 8, 9];
1682
1683 let set: TreeSet<int> = xs.iter().map(|&x| x).collect();
1684
1685 for x in xs.iter() {
1686 assert!(set.contains(x));
1687 }
1688 }
1689 }
libcollections/treemap.rs:453:1-453:1 -fn- definition:
fn mut_deref<K, V>(x: &mut Option<Box<TreeNode<K, V>>>)
-> *mut TreeNode<K, V> {
match *x {
references:- 6340: };
341: self.node = $deref(next_node);
342: }
--
377: let node = unsafe {addr!(& $($addr_mut)* *self.node)};
378: self.node = $deref(addr!(& $($addr_mut)* node.left));
379: self.stack.push(node);
--
384: let node = unsafe {addr!(& $($addr_mut)* *self.node)};
385: self.node = $deref(addr!(& $($addr_mut)* node.right));
386: }
libcollections/treemap.rs:663:38-663:38 -struct- definition:
/// Lazy backward iterator over a set
pub struct RevSetItems<'a, T> {
iter: RevEntries<'a, T, ()>
references:- 3616: #[inline]
617: pub fn rev_iter<'a>(&'a self) -> RevSetItems<'a, T> {
618: RevSetItems{iter: self.map.rev_iter()}
619: }
libcollections/treemap.rs:843:1-843:1 -fn- definition:
fn remove<K: TotalOrd, V>(node: &mut Option<Box<TreeNode<K, V>>>,
key: &K) -> Option<V> {
fn heir_swap<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>,
references:- 4877: save.left = Some(left);
878: (remove(&mut save.left, key), true)
879: } else {
libcollections/treemap.rs:680:72-680:72 -struct- definition:
/// Lazy iterator producing elements in the set intersection (in-order)
pub struct IntersectionItems<'a, T> {
a: Peekable<&'a T, SetItems<'a, T>>,
references:- 3726: impl<'a, T: TotalOrd> Iterator<&'a T> for IntersectionItems<'a, T> {
727: fn next(&mut self) -> Option<&'a T> {
libcollections/treemap.rs:813:1-813:1 -fn- definition:
fn insert<K: TotalOrd, V>(node: &mut Option<Box<TreeNode<K, V>>>,
key: K, value: V) -> Option<V> {
match *node {
references:- 3819: Less => {
820: let inserted = insert(&mut save.left, key, value);
821: skew(save);
--
825: Greater => {
826: let inserted = insert(&mut save.right, key, value);
827: skew(save);
libcollections/treemap.rs:443:1-443:1 -fn- definition:
fn deref<'a, K, V>(node: &'a Option<Box<TreeNode<K, V>>>) -> *TreeNode<K, V> {
match *node {
Some(ref n) => {
references:- 6340: };
341: self.node = $deref(next_node);
342: }
--
384: let node = unsafe {addr!(& $($addr_mut)* *self.node)};
385: self.node = $deref(addr!(& $($addr_mut)* node.right));
386: }
libcollections/treemap.rs:776:49-776:49 -fn- definition:
// Remove left horizontal link by rotating right
fn skew<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>) {
if node.left.as_ref().map_or(false, |x| x.level == node.level) {
references:- 5820: let inserted = insert(&mut save.left, key, value);
821: skew(save);
822: split(save);
--
909: for right in save.right.mut_iter() {
910: skew(right);
911: for x in right.right.mut_iter() { skew(x) }
912: }
libcollections/treemap.rs:686:72-686:72 -struct- definition:
/// Lazy iterator producing elements in the set intersection (in-order)
pub struct UnionItems<'a, T> {
a: Peekable<&'a T, SetItems<'a, T>>,
references:- 3653: pub fn union<'a>(&'a self, other: &'a TreeSet<T>) -> UnionItems<'a, T> {
654: UnionItems{a: self.iter().peekable(), b: other.iter().peekable()}
655: }
--
744: impl<'a, T: TotalOrd> Iterator<&'a T> for UnionItems<'a, T> {
745: fn next(&mut self) -> Option<&'a T> {
libcollections/treemap.rs:798:1-798:1 -fn- definition:
fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<Box<TreeNode<K, V>>>,
key: &K)
-> Option<&'r mut V> {
references:- 3804: match key.cmp(&x.key) {
805: Less => find_mut(&mut x.left, key),
806: Greater => find_mut(&mut x.right, key),
807: Equal => Some(&mut x.value),
libcollections/treemap.rs:692:81-692:81 -fn- definition:
/// Compare `x` and `y`, but return `short` if x is None and `long` if y is None
fn cmp_opt<T: TotalOrd>(x: Option<&T>, y: Option<&T>,
short: Ordering, long: Ordering) -> Ordering {
references:- 3716: loop {
717: match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) {
718: Less => return self.a.next(),
--
746: loop {
747: match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) {
748: Less => return self.a.next(),
libcollections/treemap.rs:787:14-787:14 -fn- definition:
// the parent
fn split<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>) {
if node.right.as_ref().map_or(false,
references:- 4821: skew(save);
822: split(save);
823: inserted
--
914: split(save);
915: for x in save.right.mut_iter() { split(x) }
916: }
libcollections/treemap.rs:268:38-268:38 -struct- definition:
/// Lazy backward iterator over a map
pub struct RevEntries<'a, K, V> {
iter: Entries<'a, K, V>,
references:- 4414: // the reverse Iterator impl.
415: item!(impl<'a, K, V> Iterator<(&'a K, &'a $($addr_mut)* V)> for $rev_name<'a, K, V> {
416: fn next(&mut self) -> Option<(&'a K, &'a $($addr_mut)* V)> {
--
664: pub struct RevSetItems<'a, T> {
665: iter: RevEntries<'a, T, ()>
666: }
libcollections/treemap.rs:759:19-759:19 -struct- definition:
struct TreeNode<K, V> {
key: K,
value: V,
references:- 36libcollections/treemap.rs:274:16-274:16 -struct- definition:
/// the values.
pub struct MutEntries<'a, K, V> {
stack: Vec<&'a mut TreeNode<K, V>>,
references:- 9229: fn mut_iter_for_traversal<'a>(&'a mut self) -> MutEntries<'a, K, V> {
230: MutEntries {
231: stack: vec!(),
--
399: // the forward Iterator impl.
400: item!(impl<'a, K, V> Iterator<(&'a K, &'a $($addr_mut)* V)> for $name<'a, K, V> {
401: /// Advance the iterator to the next node (in order) and return a
libcollections/treemap.rs:658:37-658:37 -struct- definition:
/// Lazy forward iterator over a set
pub struct SetItems<'a, T> {
iter: Entries<'a, T, ()>
references:- 15631: pub fn upper_bound<'a>(&'a self, v: &T) -> SetItems<'a, T> {
632: SetItems{iter: self.map.upper_bound(v)}
633: }
--
676: a: Peekable<&'a T, SetItems<'a, T>>,
677: b: Peekable<&'a T, SetItems<'a, T>>,
678: }
--
688: a: Peekable<&'a T, SetItems<'a, T>>,
689: b: Peekable<&'a T, SetItems<'a, T>>,
690: }
libcollections/treemap.rs:668:70-668:70 -struct- definition:
/// Lazy iterator producing elements in the set difference (in-order)
pub struct DifferenceItems<'a, T> {
a: Peekable<&'a T, SetItems<'a, T>>,
references:- 3636: pub fn difference<'a>(&'a self, other: &'a TreeSet<T>) -> DifferenceItems<'a, T> {
637: DifferenceItems{a: self.iter().peekable(), b: other.iter().peekable()}
638: }
--
702: impl<'a, T: TotalOrd> Iterator<&'a T> for DifferenceItems<'a, T> {
703: fn next(&mut self) -> Option<&'a T> {
libcollections/treemap.rs:674:80-674:80 -struct- definition:
/// Lazy iterator producing elements in the set symmetric difference (in-order)
pub struct SymDifferenceItems<'a, T> {
a: Peekable<&'a T, SetItems<'a, T>>,
references:- 3642: -> SymDifferenceItems<'a, T> {
643: SymDifferenceItems{a: self.iter().peekable(), b: other.iter().peekable()}
644: }
--
714: impl<'a, T: TotalOrd> Iterator<&'a T> for SymDifferenceItems<'a, T> {
715: fn next(&mut self) -> Option<&'a T> {
libcollections/treemap.rs:467:75-467:75 -struct- definition:
/// Lazy forward iterator over a map that consumes the map while iterating
pub struct MoveEntries<K, V> {
stack: Vec<TreeNode<K, V>>,
references:- 3161: };
162: MoveEntries {
163: stack: stk,
--
473: impl<K, V> Iterator<(K, V)> for MoveEntries<K,V> {
474: #[inline]
libcollections/treemap.rs:846:4-846:4 -fn- definition:
fn heir_swap<K: TotalOrd, V>(node: &mut Box<TreeNode<K, V>>,
child: &mut Option<Box<TreeNode<K, V>>>) {
// *could* be done without recursion, but it won't borrow check
references:- 2871: if left.right.is_some() {
872: heir_swap(save, &mut left.right);
873: } else {
libcollections/treemap.rs:37:19-37:19 -struct- definition:
pub struct TreeMap<K, V> {
root: Option<Box<TreeNode<K, V>>>,
length: uint
references:- 23118: /// Create an empty TreeMap
119: pub fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} }
--
203: impl<K: TotalOrd, V> TreeMap<K, V> {
204: /// Get a lazy iterator that should be initialized using
--
535: pub struct TreeSet<T> {
536: map: TreeMap<T, ()>
537: }
--
927: impl<K: TotalOrd, V> FromIterator<(K, V)> for TreeMap<K, V> {
928: fn from_iter<T: Iterator<(K, V)>>(iter: T) -> TreeMap<K, V> {
929: let mut map = TreeMap::new();
--
935: impl<K: TotalOrd, V> Extendable<(K, V)> for TreeMap<K, V> {
936: #[inline]
libcollections/treemap.rs:257:37-257:37 -struct- definition:
/// Lazy forward iterator over a map
pub struct Entries<'a, K, V> {
stack: Vec<&'a TreeNode<K, V>>,
references:- 10123: pub fn iter<'a>(&'a self) -> Entries<'a, K, V> {
124: Entries {
125: stack: vec!(),
--
206: fn iter_for_traversal<'a>(&'a self) -> Entries<'a, K, V> {
207: Entries {
208: stack: vec!(),
--
399: // the forward Iterator impl.
400: item!(impl<'a, K, V> Iterator<(&'a K, &'a $($addr_mut)* V)> for $name<'a, K, V> {
401: /// Advance the iterator to the next node (in order) and return a
--
659: pub struct SetItems<'a, T> {
660: iter: Entries<'a, T, ()>
661: }
libcollections/treemap.rs:534:19-534:19 -struct- definition:
pub struct TreeSet<T> {
map: TreeMap<T, ()>
}
references:- 24604: #[inline]
605: pub fn new() -> TreeSet<T> { TreeSet{map: TreeMap::new()} }
--
944: impl<T: TotalOrd> FromIterator<T> for TreeSet<T> {
945: fn from_iter<Iter: Iterator<T>>(iter: Iter) -> TreeSet<T> {
946: let mut set = TreeSet::new();
--
952: impl<T: TotalOrd> Extendable<T> for TreeSet<T> {
953: #[inline]
libcollections/treemap.rs:301:38-301:38 -struct- definition:
/// Lazy backward iterator over a map
pub struct RevMutEntries<'a, K, V> {
iter: MutEntries<'a, K, V>,
references:- 3414: // the reverse Iterator impl.
415: item!(impl<'a, K, V> Iterator<(&'a K, &'a $($addr_mut)* V)> for $rev_name<'a, K, V> {
416: fn next(&mut self) -> Option<(&'a K, &'a $($addr_mut)* V)> {