src/libextra/treemap.rs

   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
  16 use std::num;
  17 use std::util::{swap, replace};
  18 use std::iterator::{FromIterator, Extendable};
  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     priv root: Option<~TreeNode<K, V>>,
  40     priv 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     #[inline]
  68     fn le(&self, other: &TreeMap<K, V>) -> bool { !lt(other, self) }
  69     #[inline]
  70     fn ge(&self, other: &TreeMap<K, V>) -> bool { !lt(self, other) }
  71     #[inline]
  72     fn gt(&self, other: &TreeMap<K, V>) -> bool { lt(other, self) }
  73 }
  74
  75 impl<K: TotalOrd, V> Container for TreeMap<K, V> {
  76     /// Return the number of elements in the map
  77     fn len(&self) -> uint { self.length }
  78
  79     /// Return true if the map contains no elements
  80     fn is_empty(&self) -> bool { self.root.is_none() }
  81 }
  82
  83 impl<K: TotalOrd, V> Mutable for TreeMap<K, V> {
  84     /// Clear the map, removing all key-value pairs.
  85     fn clear(&mut self) {
  86         self.root = None;
  87         self.length = 0
  88     }
  89 }
  90
  91 impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
  92     /// Return a reference to the value corresponding to the key
  93     fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
  94         let mut current: &'a Option<~TreeNode<K, V>> = &self.root;
  95         loop {
  96             match *current {
  97               Some(ref r) => {
  98                 match key.cmp(&r.key) {
  99                   Less => current = &r.left,
 100                   Greater => current = &r.right,
 101                   Equal => return Some(&r.value)
 102                 }
 103               }
 104               None => return None
 105             }
 106         }
 107     }
 108 }
 109
 110 impl<K: TotalOrd, V> MutableMap<K, V> for TreeMap<K, V> {
 111     /// Return a mutable reference to the value corresponding to the key
 112     #[inline]
 113     fn find_mut<'a>(&'a mut self, key: &K) -> Option<&'a mut V> {
 114         find_mut(&mut self.root, key)
 115     }
 116
 117     /// Insert a key-value pair from the map. If the key already had a value
 118     /// present in the map, that value is returned. Otherwise None is returned.
 119     fn swap(&mut self, key: K, value: V) -> Option<V> {
 120         let ret = insert(&mut self.root, key, value);
 121         if ret.is_none() { self.length += 1 }
 122         ret
 123     }
 124
 125     /// Removes a key from the map, returning the value at the key if the key
 126     /// was previously in the map.
 127     fn pop(&mut self, key: &K) -> Option<V> {
 128         let ret = remove(&mut self.root, key);
 129         if ret.is_some() { self.length -= 1 }
 130         ret
 131     }
 132 }
 133
 134 impl<K: TotalOrd, V> TreeMap<K, V> {
 135     /// Create an empty TreeMap
 136     pub fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} }
 137
 138     /// Visit all keys in order
 139     pub fn each_key(&self, f: &fn(&K) -> bool) -> bool {
 140         self.iter().advance(|(k, _)| f(k))
 141     }
 142
 143     /// Visit all values in order
 144     pub fn each_value<'a>(&'a self, f: &fn(&'a V) -> bool) -> bool {
 145         self.iter().advance(|(_, v)| f(v))
 146     }
 147
 148     /// Iterate over the map and mutate the contained values
 149     pub fn mutate_values(&mut self, f: &fn(&K, &mut V) -> bool) -> bool {
 150         mutate_values(&mut self.root, f)
 151     }
 152
 153     /// Visit all key-value pairs in reverse order
 154     pub fn each_reverse<'a>(&'a self, f: &fn(&'a K, &'a V) -> bool) -> bool {
 155         each_reverse(&self.root, f)
 156     }
 157
 158     /// Visit all keys in reverse order
 159     pub fn each_key_reverse(&self, f: &fn(&K) -> bool) -> bool {
 160         self.each_reverse(|k, _| f(k))
 161     }
 162
 163     /// Visit all values in reverse order
 164     pub fn each_value_reverse(&self, f: &fn(&V) -> bool) -> bool {
 165         self.each_reverse(|_, v| f(v))
 166     }
 167
 168     /// Get a lazy iterator over the key-value pairs in the map.
 169     /// Requires that it be frozen (immutable).
 170     pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
 171         TreeMapIterator {
 172             stack: ~[],
 173             node: &self.root,
 174             remaining_min: self.length,
 175             remaining_max: self.length
 176         }
 177     }
 178
 179     /// Get a lazy iterator that should be initialized using
 180     /// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`.
 181     fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
 182         TreeMapIterator {
 183             stack: ~[],
 184             node: &self.root,
 185             remaining_min: 0,
 186             remaining_max: self.length
 187         }
 188     }
 189
 190     /// Return a lazy iterator to the first key-value pair whose key is not less than `k`
 191     /// If all keys in map are less than `k` an empty iterator is returned.
 192     pub fn lower_bound_iter<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> {
 193         let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal();
 194         loop {
 195             match *iter.node {
 196               Some(ref r) => {
 197                 match k.cmp(&r.key) {
 198                   Less => iter_traverse_left(&mut iter),
 199                   Greater => iter_traverse_right(&mut iter),
 200                   Equal => {
 201                     iter_traverse_complete(&mut iter);
 202                     return iter;
 203                   }
 204                 }
 205               }
 206               None => {
 207                 iter_traverse_complete(&mut iter);
 208                 return iter;
 209               }
 210             }
 211         }
 212     }
 213
 214     /// Return a lazy iterator to the first key-value pair whose key is greater than `k`
 215     /// If all keys in map are not greater than `k` an empty iterator is returned.
 216     pub fn upper_bound_iter<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> {
 217         let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal();
 218         loop {
 219             match *iter.node {
 220               Some(ref r) => {
 221                 match k.cmp(&r.key) {
 222                   Less => iter_traverse_left(&mut iter),
 223                   Greater => iter_traverse_right(&mut iter),
 224                   Equal => iter_traverse_right(&mut iter)
 225                 }
 226               }
 227               None => {
 228                 iter_traverse_complete(&mut iter);
 229                 return iter;
 230               }
 231             }
 232         }
 233     }
 234
 235     /// Get a lazy iterator that consumes the treemap.
 236     pub fn consume_iter(self) -> TreeMapConsumeIterator<K, V> {
 237         let TreeMap { root: root, length: length } = self;
 238         let stk = match root {
 239             None => ~[],
 240             Some(~tn) => ~[tn]
 241         };
 242         TreeMapConsumeIterator {
 243             stack: stk,
 244             remaining: length
 245         }
 246     }
 247 }
 248
 249 /// Lazy forward iterator over a map
 250 pub struct TreeMapIterator<'self, K, V> {
 251     priv stack: ~[&'self ~TreeNode<K, V>],
 252     priv node: &'self Option<~TreeNode<K, V>>,
 253     priv remaining_min: uint,
 254     priv remaining_max: uint
 255 }
 256
 257 impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V> {
 258     /// Advance the iterator to the next node (in order) and return a
 259     /// tuple with a reference to the key and value. If there are no
 260     /// more nodes, return `None`.
 261     fn next(&mut self) -> Option<(&'self K, &'self V)> {
 262         while !self.stack.is_empty() || self.node.is_some() {
 263             match *self.node {
 264               Some(ref x) => {
 265                 self.stack.push(x);
 266                 self.node = &x.left;
 267               }
 268               None => {
 269                 let res = self.stack.pop();
 270                 self.node = &res.right;
 271                 self.remaining_max -= 1;
 272                 if self.remaining_min > 0 {
 273                     self.remaining_min -= 1;
 274                 }
 275                 return Some((&res.key, &res.value));
 276               }
 277             }
 278         }
 279         None
 280     }
 281
 282     #[inline]
 283     fn size_hint(&self) -> (uint, Option<uint>) {
 284         (self.remaining_min, Some(self.remaining_max))
 285     }
 286 }
 287
 288 /// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
 289 /// initialize TreeMapIterator pointing to element inside tree structure.
 290 ///
 291 /// They should be used in following manner:
 292 ///   - create iterator using TreeMap::iter_for_traversal
 293 ///   - find required node using `iter_traverse_left`/`iter_traverse_right`
 294 ///     (current node is `TreeMapIterator::node` field)
 295 ///   - complete initialization with `iter_traverse_complete`
 296 #[inline]
 297 fn iter_traverse_left<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
 298     let node = it.node.get_ref();
 299     it.stack.push(node);
 300     it.node = &node.left;
 301 }
 302
 303 #[inline]
 304 fn iter_traverse_right<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
 305     it.node = &(it.node.get_ref().right);
 306 }
 307
 308 /// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
 309 /// initialize TreeMapIterator pointing to element inside tree structure.
 310 ///
 311 /// Completes traversal. Should be called before using iterator.
 312 /// Iteration will start from `self.node`.
 313 /// If `self.node` is None iteration will start from last node from which we
 314 /// traversed left.
 315 #[inline]
 316 fn iter_traverse_complete<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
 317     static none: Option<~TreeNode<K, V>> = None;
 318     match *it.node {
 319         Some(ref n) => {
 320             it.stack.push(n);
 321             it.node = &none;
 322         }
 323         None => ()
 324     }
 325 }
 326
 327 /// Lazy forward iterator over a map that consumes the map while iterating
 328 pub struct TreeMapConsumeIterator<K, V> {
 329     priv stack: ~[TreeNode<K, V>],
 330     priv remaining: uint
 331 }
 332
 333 impl<K, V> Iterator<(K, V)> for TreeMapConsumeIterator<K,V> {
 334     #[inline]
 335     fn next(&mut self) -> Option<(K, V)> {
 336         while !self.stack.is_empty() {
 337             let TreeNode {
 338                 key: key,
 339                 value: value,
 340                 left: left,
 341                 right: right,
 342                 level: level
 343             } = self.stack.pop();
 344
 345             match left {
 346                 Some(~left) => {
 347                     let n = TreeNode {
 348                         key: key,
 349                         value: value,
 350                         left: None,
 351                         right: right,
 352                         level: level
 353                     };
 354                     self.stack.push(n);
 355                     self.stack.push(left);
 356                 }
 357                 None => {
 358                     match right {
 359                         Some(~right) => self.stack.push(right),
 360                         None => ()
 361                     }
 362                     self.remaining -= 1;
 363                     return Some((key, value))
 364                 }
 365             }
 366         }
 367         None
 368     }
 369
 370     #[inline]
 371     fn size_hint(&self) -> (uint, Option<uint>) {
 372         (self.remaining, Some(self.remaining))
 373     }
 374
 375 }
 376
 377 impl<'self, T> Iterator<&'self T> for TreeSetIterator<'self, T> {
 378     /// Advance the iterator to the next node (in order). If there are no more nodes, return `None`.
 379     #[inline]
 380     fn next(&mut self) -> Option<&'self T> {
 381         do self.iter.next().map_move |(value, _)| { value }
 382     }
 383 }
 384
 385 /// A implementation of the `Set` trait on top of the `TreeMap` container. The
 386 /// only requirement is that the type of the elements contained ascribes to the
 387 /// `TotalOrd` trait.
 388 pub struct TreeSet<T> {
 389     priv map: TreeMap<T, ()>
 390 }
 391
 392 impl<T: Eq + TotalOrd> Eq for TreeSet<T> {
 393     #[inline]
 394     fn eq(&self, other: &TreeSet<T>) -> bool { self.map == other.map }
 395     #[inline]
 396     fn ne(&self, other: &TreeSet<T>) -> bool { self.map != other.map }
 397 }
 398
 399 impl<T: Ord + TotalOrd> Ord for TreeSet<T> {
 400     #[inline]
 401     fn lt(&self, other: &TreeSet<T>) -> bool { self.map < other.map }
 402     #[inline]
 403     fn le(&self, other: &TreeSet<T>) -> bool { self.map <= other.map }
 404     #[inline]
 405     fn ge(&self, other: &TreeSet<T>) -> bool { self.map >= other.map }
 406     #[inline]
 407     fn gt(&self, other: &TreeSet<T>) -> bool { self.map > other.map }
 408 }
 409
 410 impl<T: TotalOrd> Container for TreeSet<T> {
 411     /// Return the number of elements in the set
 412     #[inline]
 413     fn len(&self) -> uint { self.map.len() }
 414
 415     /// Return true if the set contains no elements
 416     #[inline]
 417     fn is_empty(&self) -> bool { self.map.is_empty() }
 418 }
 419
 420 impl<T: TotalOrd> Mutable for TreeSet<T> {
 421     /// Clear the set, removing all values.
 422     #[inline]
 423     fn clear(&mut self) { self.map.clear() }
 424 }
 425
 426 impl<T: TotalOrd> Set<T> for TreeSet<T> {
 427     /// Return true if the set contains a value
 428     #[inline]
 429     fn contains(&self, value: &T) -> bool {
 430         self.map.contains_key(value)
 431     }
 432
 433     /// Return true if the set has no elements in common with `other`.
 434     /// This is equivalent to checking for an empty intersection.
 435     fn is_disjoint(&self, other: &TreeSet<T>) -> bool {
 436         self.intersection(other).next().is_none()
 437     }
 438
 439     /// Return true if the set is a subset of another
 440     #[inline]
 441     fn is_subset(&self, other: &TreeSet<T>) -> bool {
 442         other.is_superset(self)
 443     }
 444
 445     /// Return true if the set is a superset of another
 446     fn is_superset(&self, other: &TreeSet<T>) -> bool {
 447         let mut x = self.iter();
 448         let mut y = other.iter();
 449         let mut a = x.next();
 450         let mut b = y.next();
 451         while b.is_some() {
 452             if a.is_none() {
 453                 return false
 454             }
 455
 456             let a1 = a.unwrap();
 457             let b1 = b.unwrap();
 458
 459             match a1.cmp(b1) {
 460               Less => (),
 461               Greater => return false,
 462               Equal => b = y.next(),
 463             }
 464
 465             a = x.next();
 466         }
 467         true
 468     }
 469 }
 470
 471 impl<T: TotalOrd> MutableSet<T> for TreeSet<T> {
 472     /// Add a value to the set. Return true if the value was not already
 473     /// present in the set.
 474     #[inline]
 475     fn insert(&mut self, value: T) -> bool { self.map.insert(value, ()) }
 476
 477     /// Remove a value from the set. Return true if the value was
 478     /// present in the set.
 479     #[inline]
 480     fn remove(&mut self, value: &T) -> bool { self.map.remove(value) }
 481 }
 482
 483 impl<T: TotalOrd> TreeSet<T> {
 484     /// Create an empty TreeSet
 485     #[inline]
 486     pub fn new() -> TreeSet<T> { TreeSet{map: TreeMap::new()} }
 487
 488     /// Get a lazy iterator over the values in the set.
 489     /// Requires that it be frozen (immutable).
 490     #[inline]
 491     pub fn iter<'a>(&'a self) -> TreeSetIterator<'a, T> {
 492         TreeSetIterator{iter: self.map.iter()}
 493     }
 494
 495     /// Get a lazy iterator pointing to the first value not less than `v` (greater or equal).
 496     /// If all elements in the set are less than `v` empty iterator is returned.
 497     #[inline]
 498     pub fn lower_bound_iter<'a>(&'a self, v: &T) -> TreeSetIterator<'a, T> {
 499         TreeSetIterator{iter: self.map.lower_bound_iter(v)}
 500     }
 501
 502     /// Get a lazy iterator pointing to the first value greater than `v`.
 503     /// If all elements in the set are not greater than `v` empty iterator is returned.
 504     #[inline]
 505     pub fn upper_bound_iter<'a>(&'a self, v: &T) -> TreeSetIterator<'a, T> {
 506         TreeSetIterator{iter: self.map.upper_bound_iter(v)}
 507     }
 508
 509     /// Visit all values in reverse order
 510     #[inline]
 511     pub fn each_reverse(&self, f: &fn(&T) -> bool) -> bool {
 512         self.map.each_key_reverse(f)
 513     }
 514
 515     /// Visit the values (in-order) representing the difference
 516     pub fn difference<'a>(&'a self, other: &'a TreeSet<T>) -> Difference<'a, T> {
 517         Difference{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
 518     }
 519
 520     /// Visit the values (in-order) representing the symmetric difference
 521     pub fn symmetric_difference<'a>(&'a self, other: &'a TreeSet<T>)
 522         -> SymDifference<'a, T> {
 523         SymDifference{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
 524     }
 525
 526     /// Visit the values (in-order) representing the intersection
 527     pub fn intersection<'a>(&'a self, other: &'a TreeSet<T>)
 528         -> Intersection<'a, T> {
 529         Intersection{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
 530     }
 531
 532     /// Visit the values (in-order) representing the union
 533     pub fn union<'a>(&'a self, other: &'a TreeSet<T>) -> Union<'a, T> {
 534         Union{a: Focus::new(self.iter()), b: Focus::new(other.iter())}
 535     }
 536 }
 537
 538 /// Lazy forward iterator over a set
 539 pub struct TreeSetIterator<'self, T> {
 540     priv iter: TreeMapIterator<'self, T, ()>
 541 }
 542
 543 // Encapsulate an iterator and hold its latest value until stepped forward
 544 struct Focus<A, T> {
 545     priv iter: T,
 546     priv focus: Option<A>,
 547 }
 548
 549 impl<A, T: Iterator<A>> Focus<A, T> {
 550     fn new(mut it: T) -> Focus<A, T> {
 551         Focus{focus: it.next(), iter: it}
 552     }
 553     fn step(&mut self) {
 554         self.focus = self.iter.next()
 555     }
 556 }
 557
 558 /// Lazy iterator producing elements in the set difference (in-order)
 559 pub struct Difference<'self, T> {
 560     priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
 561     priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
 562 }
 563
 564 /// Lazy iterator producing elements in the set symmetric difference (in-order)
 565 pub struct SymDifference<'self, T> {
 566     priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
 567     priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
 568 }
 569
 570 /// Lazy iterator producing elements in the set intersection (in-order)
 571 pub struct Intersection<'self, T> {
 572     priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
 573     priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
 574 }
 575
 576 /// Lazy iterator producing elements in the set intersection (in-order)
 577 pub struct Union<'self, T> {
 578     priv a: Focus<&'self T, TreeSetIterator<'self, T>>,
 579     priv b: Focus<&'self T, TreeSetIterator<'self, T>>,
 580 }
 581
 582 impl<'self, T: TotalOrd> Iterator<&'self T> for Difference<'self, T> {
 583     fn next(&mut self) -> Option<&'self T> {
 584         loop {
 585             match (self.a.focus, self.b.focus) {
 586                 (None    , _       ) => return None,
 587                 (ret     , None    ) => { self.a.step(); return ret },
 588                 (Some(a1), Some(b1)) => {
 589                     let cmp = a1.cmp(b1);
 590                     if cmp == Less {
 591                         self.a.step();
 592                         return Some(a1);
 593                     } else {
 594                         if cmp == Equal { self.a.step() }
 595                         self.b.step();
 596                     }
 597                 }
 598             }
 599         }
 600     }
 601 }
 602
 603 impl<'self, T: TotalOrd> Iterator<&'self T> for SymDifference<'self, T> {
 604     fn next(&mut self) -> Option<&'self T> {
 605         loop {
 606             match (self.a.focus, self.b.focus) {
 607                 (ret     , None    ) => { self.a.step(); return ret },
 608                 (None    , ret     ) => { self.b.step(); return ret },
 609                 (Some(a1), Some(b1)) => {
 610                     let cmp = a1.cmp(b1);
 611                     if cmp == Less {
 612                         self.a.step();
 613                         return Some(a1);
 614                     } else {
 615                         self.b.step();
 616                         if cmp == Greater {
 617                             return Some(b1);
 618                         } else {
 619                             self.a.step();
 620                         }
 621                     }
 622                 }
 623             }
 624         }
 625     }
 626 }
 627
 628 impl<'self, T: TotalOrd> Iterator<&'self T> for Intersection<'self, T> {
 629     fn next(&mut self) -> Option<&'self T> {
 630         loop {
 631             match (self.a.focus, self.b.focus) {
 632                 (None    , _       ) => return None,
 633                 (_       , None    ) => return None,
 634                 (Some(a1), Some(b1)) => {
 635                     let cmp = a1.cmp(b1);
 636                     if cmp == Less {
 637                         self.a.step();
 638                     } else {
 639                         self.b.step();
 640                         if cmp == Equal {
 641                             return Some(a1);
 642                         }
 643                     }
 644                 },
 645             }
 646         }
 647     }
 648 }
 649
 650 impl<'self, T: TotalOrd> Iterator<&'self T> for Union<'self, T> {
 651     fn next(&mut self) -> Option<&'self T> {
 652         loop {
 653             match (self.a.focus, self.b.focus) {
 654                 (ret     , None) => { self.a.step(); return ret },
 655                 (None    , ret ) => { self.b.step(); return ret },
 656                 (Some(a1), Some(b1)) => {
 657                     let cmp = a1.cmp(b1);
 658                     if cmp == Greater {
 659                         self.b.step();
 660                         return Some(b1);
 661                     } else {
 662                         self.a.step();
 663                         if cmp == Equal {
 664                             self.b.step();
 665                         }
 666                         return Some(a1);
 667                     }
 668                 }
 669             }
 670         }
 671     }
 672 }
 673
 674
 675 // Nodes keep track of their level in the tree, starting at 1 in the
 676 // leaves and with a red child sharing the level of the parent.
 677 #[deriving(Clone)]
 678 struct TreeNode<K, V> {
 679     key: K,
 680     value: V,
 681     left: Option<~TreeNode<K, V>>,
 682     right: Option<~TreeNode<K, V>>,
 683     level: uint
 684 }
 685
 686 impl<K: TotalOrd, V> TreeNode<K, V> {
 687     /// Creates a new tree node.
 688     #[inline]
 689     pub fn new(key: K, value: V) -> TreeNode<K, V> {
 690         TreeNode{key: key, value: value, left: None, right: None, level: 1}
 691     }
 692 }
 693
 694 fn each<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>,
 695                             f: &fn(&'r K, &'r V) -> bool) -> bool {
 696     node.iter().advance(|x| each(&x.left,  |k,v| f(k,v)) && f(&x.key, &x.value) &&
 697                             each(&x.right, |k,v| f(k,v)))
 698 }
 699
 700 fn each_reverse<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>,
 701                                     f: &fn(&'r K, &'r V) -> bool) -> bool {
 702     node.iter().advance(|x| each_reverse(&x.right, |k,v| f(k,v)) && f(&x.key, &x.value) &&
 703                             each_reverse(&x.left,  |k,v| f(k,v)))
 704 }
 705
 706 fn mutate_values<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
 707                                      f: &fn(&'r K, &'r mut V) -> bool)
 708                                   -> bool {
 709     match *node {
 710       Some(~TreeNode{key: ref key, value: ref mut value, left: ref mut left,
 711                      right: ref mut right, _}) => {
 712         if !mutate_values(left,  |k,v| f(k,v)) { return false }
 713         if !f(key, value) { return false }
 714         if !mutate_values(right, |k,v| f(k,v)) { return false }
 715       }
 716       None => return false
 717     }
 718     true
 719 }
 720
 721 // Remove left horizontal link by rotating right
 722 fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
 723     if node.left.map_default(false, |x| x.level == node.level) {
 724         let mut save = node.left.take_unwrap();
 725         swap(&mut node.left, &mut save.right); // save.right now None
 726         swap(node, &mut save);
 727         node.right = Some(save);
 728     }
 729 }
 730
 731 // Remove dual horizontal link by rotating left and increasing level of
 732 // the parent
 733 fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
 734     if node.right.map_default(false,
 735       |x| x.right.map_default(false, |y| y.level == node.level)) {
 736         let mut save = node.right.take_unwrap();
 737         swap(&mut node.right, &mut save.left); // save.left now None
 738         save.level += 1;
 739         swap(node, &mut save);
 740         node.left = Some(save);
 741     }
 742 }
 743
 744 fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
 745                                 key: &K)
 746                              -> Option<&'r mut V> {
 747     match *node {
 748       Some(ref mut x) => {
 749         match key.cmp(&x.key) {
 750           Less => find_mut(&mut x.left, key),
 751           Greater => find_mut(&mut x.right, key),
 752           Equal => Some(&mut x.value),
 753         }
 754       }
 755       None => None
 756     }
 757 }
 758
 759 fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
 760                           key: K, value: V) -> Option<V> {
 761     match *node {
 762       Some(ref mut save) => {
 763         match key.cmp(&save.key) {
 764           Less => {
 765             let inserted = insert(&mut save.left, key, value);
 766             skew(save);
 767             split(save);
 768             inserted
 769           }
 770           Greater => {
 771             let inserted = insert(&mut save.right, key, value);
 772             skew(save);
 773             split(save);
 774             inserted
 775           }
 776           Equal => {
 777             save.key = key;
 778             Some(replace(&mut save.value, value))
 779           }
 780         }
 781       }
 782       None => {
 783        *node = Some(~TreeNode::new(key, value));
 784         None
 785       }
 786     }
 787 }
 788
 789 fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
 790                           key: &K) -> Option<V> {
 791     fn heir_swap<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>,
 792                                  child: &mut Option<~TreeNode<K, V>>) {
 793         // *could* be done without recursion, but it won't borrow check
 794         for x in child.mut_iter() {
 795             if x.right.is_some() {
 796                 heir_swap(node, &mut x.right);
 797             } else {
 798                 swap(&mut node.key, &mut x.key);
 799                 swap(&mut node.value, &mut x.value);
 800             }
 801         }
 802     }
 803
 804     match *node {
 805       None => {
 806         return None; // bottom of tree
 807       }
 808       Some(ref mut save) => {
 809         let (ret, rebalance) = match key.cmp(&save.key) {
 810           Less => (remove(&mut save.left, key), true),
 811           Greater => (remove(&mut save.right, key), true),
 812           Equal => {
 813             if save.left.is_some() {
 814                 if save.right.is_some() {
 815                     let mut left = save.left.take_unwrap();
 816                     if left.right.is_some() {
 817                         heir_swap(save, &mut left.right);
 818                     } else {
 819                         swap(&mut save.key, &mut left.key);
 820                         swap(&mut save.value, &mut left.value);
 821                     }
 822                     save.left = Some(left);
 823                     (remove(&mut save.left, key), true)
 824                 } else {
 825                     let new = save.left.take_unwrap();
 826                     let ~TreeNode{value, _} = replace(save, new);
 827                     *save = save.left.take_unwrap();
 828                     (Some(value), true)
 829                 }
 830             } else if save.right.is_some() {
 831                 let new = save.right.take_unwrap();
 832                 let ~TreeNode{value, _} = replace(save, new);
 833                 (Some(value), true)
 834             } else {
 835                 (None, false)
 836             }
 837           }
 838         };
 839
 840         if rebalance {
 841             let left_level = save.left.map_default(0, |x| x.level);
 842             let right_level = save.right.map_default(0, |x| x.level);
 843
 844             // re-balance, if necessary
 845             if left_level < save.level - 1 || right_level < save.level - 1 {
 846                 save.level -= 1;
 847
 848                 if right_level > save.level {
 849                     for x in save.right.mut_iter() { x.level = save.level }
 850                 }
 851
 852                 skew(save);
 853
 854                 for right in save.right.mut_iter() {
 855                     skew(right);
 856                     for x in right.right.mut_iter() { skew(x) }
 857                 }
 858
 859                 split(save);
 860                 for x in save.right.mut_iter() { split(x) }
 861             }
 862
 863             return ret;
 864         }
 865       }
 866     }
 867     return match node.take() {
 868         Some(~TreeNode{value, _}) => Some(value), None => fail!()
 869     };
 870 }
 871
 872 impl<K: TotalOrd, V, T: Iterator<(K, V)>> FromIterator<(K, V), T> for TreeMap<K, V> {
 873     fn from_iterator(iter: &mut T) -> TreeMap<K, V> {
 874         let mut map = TreeMap::new();
 875         map.extend(iter);
 876         map
 877     }
 878 }
 879
 880 impl<K: TotalOrd, V, T: Iterator<(K, V)>> Extendable<(K, V), T> for TreeMap<K, V> {
 881     #[inline]
 882     fn extend(&mut self, iter: &mut T) {
 883         for (k, v) in *iter {
 884             self.insert(k, v);
 885         }
 886     }
 887 }
 888
 889 impl<T: TotalOrd, Iter: Iterator<T>> FromIterator<T, Iter> for TreeSet<T> {
 890     pub fn from_iterator(iter: &mut Iter) -> TreeSet<T> {
 891         let mut set = TreeSet::new();
 892         set.extend(iter);
 893         set
 894     }
 895 }
 896
 897 impl<T: TotalOrd, Iter: Iterator<T>> Extendable<T, Iter> for TreeSet<T> {
 898     #[inline]
 899     fn extend(&mut self, iter: &mut Iter) {
 900         for elem in *iter {
 901             self.insert(elem);
 902         }
 903     }
 904 }
 905
 906 #[cfg(test)]
 907 mod test_treemap {
 908
 909     use super::*;
 910
 911     use std::rand::RngUtil;
 912     use std::rand;
 913
 914     #[test]
 915     fn find_empty() {
 916         let m = TreeMap::new::<int, int>(); assert!(m.find(&5) == None);
 917     }
 918
 919     #[test]
 920     fn find_not_found() {
 921         let mut m = TreeMap::new();
 922         assert!(m.insert(1, 2));
 923         assert!(m.insert(5, 3));
 924         assert!(m.insert(9, 3));
 925         assert_eq!(m.find(&2), None);
 926     }
 927
 928     #[test]
 929     fn test_find_mut() {
 930         let mut m = TreeMap::new();
 931         assert!(m.insert(1, 12));
 932         assert!(m.insert(2, 8));
 933         assert!(m.insert(5, 14));
 934         let new = 100;
 935         match m.find_mut(&5) {
 936           None => fail!(), Some(x) => *x = new
 937         }
 938         assert_eq!(m.find(&5), Some(&new));
 939     }
 940
 941     #[test]
 942     fn insert_replace() {
 943         let mut m = TreeMap::new();
 944         assert!(m.insert(5, 2));
 945         assert!(m.insert(2, 9));
 946         assert!(!m.insert(2, 11));
 947         assert_eq!(m.find(&2).unwrap(), &11);
 948     }
 949
 950     #[test]
 951     fn test_clear() {
 952         let mut m = TreeMap::new();
 953         m.clear();
 954         assert!(m.insert(5, 11));
 955         assert!(m.insert(12, -3));
 956         assert!(m.insert(19, 2));
 957         m.clear();
 958         assert!(m.find(&5).is_none());
 959         assert!(m.find(&12).is_none());
 960         assert!(m.find(&19).is_none());
 961         assert!(m.is_empty());
 962     }
 963
 964     #[test]
 965     fn u8_map() {
 966         let mut m = TreeMap::new();
 967
 968         let k1 = "foo".as_bytes();
 969         let k2 = "bar".as_bytes();
 970         let v1 = "baz".as_bytes();
 971         let v2 = "foobar".as_bytes();
 972
 973         m.insert(k1.clone(), v1.clone());
 974         m.insert(k2.clone(), v2.clone());
 975
 976         assert_eq!(m.find(&k2), Some(&v2));
 977         assert_eq!(m.find(&k1), Some(&v1));
 978     }
 979
 980     fn check_equal<K: Eq + TotalOrd, V: Eq>(ctrl: &[(K, V)],
 981                                             map: &TreeMap<K, V>) {
 982         assert_eq!(ctrl.is_empty(), map.is_empty());
 983         for x in ctrl.iter() {
 984             let &(ref k, ref v) = x;
 985             assert!(map.find(k).unwrap() == v)
 986         }
 987         for (map_k, map_v) in map.iter() {
 988             let mut found = false;
 989             for x in ctrl.iter() {
 990                 let &(ref ctrl_k, ref ctrl_v) = x;
 991                 if *map_k == *ctrl_k {
 992                     assert!(*map_v == *ctrl_v);
 993                     found = true;
 994                     break;
 995                 }
 996             }
 997             assert!(found);
 998         }
 999     }
1000
1001     fn check_left<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>,
1002                                   parent: &~TreeNode<K, V>) {
1003         match *node {
1004           Some(ref r) => {
1005             assert_eq!(r.key.cmp(&parent.key), Less);
1006             assert!(r.level == parent.level - 1); // left is black
1007             check_left(&r.left, r);
1008             check_right(&r.right, r, false);
1009           }
1010           None => assert!(parent.level == 1) // parent is leaf
1011         }
1012     }
1013
1014     fn check_right<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>,
1015                                    parent: &~TreeNode<K, V>,
1016                                    parent_red: bool) {
1017         match *node {
1018           Some(ref r) => {
1019             assert_eq!(r.key.cmp(&parent.key), Greater);
1020             let red = r.level == parent.level;
1021             if parent_red { assert!(!red) } // no dual horizontal links
1022             // Right red or black
1023             assert!(red || r.level == parent.level - 1);
1024             check_left(&r.left, r);
1025             check_right(&r.right, r, red);
1026           }
1027           None => assert!(parent.level == 1) // parent is leaf
1028         }
1029     }
1030
1031     fn check_structure<K: TotalOrd, V>(map: &TreeMap<K, V>) {
1032         match map.root {
1033           Some(ref r) => {
1034             check_left(&r.left, r);
1035             check_right(&r.right, r, false);
1036           }
1037           None => ()
1038         }
1039     }
1040
1041     #[test]
1042     fn test_rand_int() {
1043         let mut map = TreeMap::new::<int, int>();
1044         let mut ctrl = ~[];
1045
1046         check_equal(ctrl, &map);
1047         assert!(map.find(&5).is_none());
1048
1049         let mut rng = rand::IsaacRng::new_seeded(&[42]);
1050
1051         do 3.times {
1052             do 90.times {
1053                 let k = rng.gen();
1054                 let v = rng.gen();
1055                 if !ctrl.iter().any(|x| x == &(k, v)) {
1056                     assert!(map.insert(k, v));
1057                     ctrl.push((k, v));
1058                     check_structure(&map);
1059                     check_equal(ctrl, &map);
1060                 }
1061             }
1062
1063             do 30.times {
1064                 let r = rng.gen_uint_range(0, ctrl.len());
1065                 let (key, _) = ctrl.remove(r);
1066                 assert!(map.remove(&key));
1067                 check_structure(&map);
1068                 check_equal(ctrl, &map);
1069             }
1070         }
1071     }
1072
1073     #[test]
1074     fn test_len() {
1075         let mut m = TreeMap::new();
1076         assert!(m.insert(3, 6));
1077         assert_eq!(m.len(), 1);
1078         assert!(m.insert(0, 0));
1079         assert_eq!(m.len(), 2);
1080         assert!(m.insert(4, 8));
1081         assert_eq!(m.len(), 3);
1082         assert!(m.remove(&3));
1083         assert_eq!(m.len(), 2);
1084         assert!(!m.remove(&5));
1085         assert_eq!(m.len(), 2);
1086         assert!(m.insert(2, 4));
1087         assert_eq!(m.len(), 3);
1088         assert!(m.insert(1, 2));
1089         assert_eq!(m.len(), 4);
1090     }
1091
1092     #[test]
1093     fn test_iterator() {
1094         let mut m = TreeMap::new();
1095
1096         assert!(m.insert(3, 6));
1097         assert!(m.insert(0, 0));
1098         assert!(m.insert(4, 8));
1099         assert!(m.insert(2, 4));
1100         assert!(m.insert(1, 2));
1101
1102         let mut n = 0;
1103         for (k, v) in m.iter() {
1104             assert_eq!(*k, n);
1105             assert_eq!(*v, n * 2);
1106             n += 1;
1107         }
1108         assert_eq!(n, 5);
1109     }
1110
1111     #[test]
1112     fn test_interval_iteration() {
1113         let mut m = TreeMap::new();
1114         for i in range(1, 100) {
1115             assert!(m.insert(i * 2, i * 4));
1116         }
1117
1118         for i in range(1, 198) {
1119             let mut lb_it = m.lower_bound_iter(&i);
1120             let (&k, &v) = lb_it.next().unwrap();
1121             let lb = i + i % 2;
1122             assert_eq!(lb, k);
1123             assert_eq!(lb * 2, v);
1124
1125             let mut ub_it = m.upper_bound_iter(&i);
1126             let (&k, &v) = ub_it.next().unwrap();
1127             let ub = i + 2 - i % 2;
1128             assert_eq!(ub, k);
1129             assert_eq!(ub * 2, v);
1130         }
1131         let mut end_it = m.lower_bound_iter(&199);
1132         assert_eq!(end_it.next(), None);
1133     }
1134
1135     #[test]
1136     fn test_each_reverse() {
1137         let mut m = TreeMap::new();
1138
1139         assert!(m.insert(3, 6));
1140         assert!(m.insert(0, 0));
1141         assert!(m.insert(4, 8));
1142         assert!(m.insert(2, 4));
1143         assert!(m.insert(1, 2));
1144
1145         let mut n = 4;
1146         do m.each_reverse |k, v| {
1147             assert_eq!(*k, n);
1148             assert_eq!(*v, n * 2);
1149             n -= 1;
1150             true
1151         };
1152     }
1153
1154     #[test]
1155     fn test_eq() {
1156         let mut a = TreeMap::new();
1157         let mut b = TreeMap::new();
1158
1159         assert!(a == b);
1160         assert!(a.insert(0, 5));
1161         assert!(a != b);
1162         assert!(b.insert(0, 4));
1163         assert!(a != b);
1164         assert!(a.insert(5, 19));
1165         assert!(a != b);
1166         assert!(!b.insert(0, 5));
1167         assert!(a != b);
1168         assert!(b.insert(5, 19));
1169         assert!(a == b);
1170     }
1171
1172     #[test]
1173     fn test_lt() {
1174         let mut a = TreeMap::new();
1175         let mut b = TreeMap::new();
1176
1177         assert!(!(a < b) && !(b < a));
1178         assert!(b.insert(0, 5));
1179         assert!(a < b);
1180         assert!(a.insert(0, 7));
1181         assert!(!(a < b) && b < a);
1182         assert!(b.insert(-2, 0));
1183         assert!(b < a);
1184         assert!(a.insert(-5, 2));
1185         assert!(a < b);
1186         assert!(a.insert(6, 2));
1187         assert!(a < b && !(b < a));
1188     }
1189
1190     #[test]
1191     fn test_ord() {
1192         let mut a = TreeMap::new();
1193         let mut b = TreeMap::new();
1194
1195         assert!(a <= b && a >= b);
1196         assert!(a.insert(1, 1));
1197         assert!(a > b && a >= b);
1198         assert!(b < a && b <= a);
1199         assert!(b.insert(2, 2));
1200         assert!(b > a && b >= a);
1201         assert!(a < b && a <= b);
1202     }
1203
1204     #[test]
1205     fn test_lazy_iterator() {
1206         let mut m = TreeMap::new();
1207         let (x1, y1) = (2, 5);
1208         let (x2, y2) = (9, 12);
1209         let (x3, y3) = (20, -3);
1210         let (x4, y4) = (29, 5);
1211         let (x5, y5) = (103, 3);
1212
1213         assert!(m.insert(x1, y1));
1214         assert!(m.insert(x2, y2));
1215         assert!(m.insert(x3, y3));
1216         assert!(m.insert(x4, y4));
1217         assert!(m.insert(x5, y5));
1218
1219         let m = m;
1220         let mut a = m.iter();
1221
1222         assert_eq!(a.next().unwrap(), (&x1, &y1));
1223         assert_eq!(a.next().unwrap(), (&x2, &y2));
1224         assert_eq!(a.next().unwrap(), (&x3, &y3));
1225         assert_eq!(a.next().unwrap(), (&x4, &y4));
1226         assert_eq!(a.next().unwrap(), (&x5, &y5));
1227
1228         assert!(a.next().is_none());
1229
1230         let mut b = m.iter();
1231
1232         let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4),
1233                         (&x5, &y5)];
1234         let mut i = 0;
1235
1236         for x in b {
1237             assert_eq!(expected[i], x);
1238             i += 1;
1239
1240             if i == 2 {
1241                 break
1242             }
1243         }
1244
1245         for x in b {
1246             assert_eq!(expected[i], x);
1247             i += 1;
1248         }
1249     }
1250
1251     #[test]
1252     fn test_from_iter() {
1253         let xs = ~[(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1254
1255         let map: TreeMap<int, int> = xs.iter().transform(|&x| x).collect();
1256
1257         for &(k, v) in xs.iter() {
1258             assert_eq!(map.find(&k), Some(&v));
1259         }
1260     }
1261
1262 }
1263
1264 #[cfg(test)]
1265 mod bench {
1266
1267     use super::*;
1268     use test::BenchHarness;
1269     use container::bench::*;
1270
1271     // Find seq
1272     #[bench]
1273     pub fn insert_rand_100(bh: &mut BenchHarness) {
1274         let mut m : TreeMap<uint,uint> = TreeMap::new();
1275         insert_rand_n(100, &mut m, bh);
1276     }
1277
1278     #[bench]
1279     pub fn insert_rand_10_000(bh: &mut BenchHarness) {
1280         let mut m : TreeMap<uint,uint> = TreeMap::new();
1281         insert_rand_n(10_000, &mut m, bh);
1282     }
1283
1284     // Insert seq
1285     #[bench]
1286     pub fn insert_seq_100(bh: &mut BenchHarness) {
1287         let mut m : TreeMap<uint,uint> = TreeMap::new();
1288         insert_seq_n(100, &mut m, bh);
1289     }
1290
1291     #[bench]
1292     pub fn insert_seq_10_000(bh: &mut BenchHarness) {
1293         let mut m : TreeMap<uint,uint> = TreeMap::new();
1294         insert_seq_n(10_000, &mut m, bh);
1295     }
1296
1297     // Find rand
1298     #[bench]
1299     pub fn find_rand_100(bh: &mut BenchHarness) {
1300         let mut m : TreeMap<uint,uint> = TreeMap::new();
1301         find_rand_n(100, &mut m, bh);
1302     }
1303
1304     #[bench]
1305     pub fn find_rand_10_000(bh: &mut BenchHarness) {
1306         let mut m : TreeMap<uint,uint> = TreeMap::new();
1307         find_rand_n(10_000, &mut m, bh);
1308     }
1309
1310     // Find seq
1311     #[bench]
1312     pub fn find_seq_100(bh: &mut BenchHarness) {
1313         let mut m : TreeMap<uint,uint> = TreeMap::new();
1314         find_seq_n(100, &mut m, bh);
1315     }
1316
1317     #[bench]
1318     pub fn find_seq_10_000(bh: &mut BenchHarness) {
1319         let mut m : TreeMap<uint,uint> = TreeMap::new();
1320         find_seq_n(10_000, &mut m, bh);
1321     }
1322 }
1323
1324 #[cfg(test)]
1325 mod test_set {
1326
1327     use super::*;
1328
1329     #[test]
1330     fn test_clear() {
1331         let mut s = TreeSet::new();
1332         s.clear();
1333         assert!(s.insert(5));
1334         assert!(s.insert(12));
1335         assert!(s.insert(19));
1336         s.clear();
1337         assert!(!s.contains(&5));
1338         assert!(!s.contains(&12));
1339         assert!(!s.contains(&19));
1340         assert!(s.is_empty());
1341     }
1342
1343     #[test]
1344     fn test_disjoint() {
1345         let mut xs = TreeSet::new();
1346         let mut ys = TreeSet::new();
1347         assert!(xs.is_disjoint(&ys));
1348         assert!(ys.is_disjoint(&xs));
1349         assert!(xs.insert(5));
1350         assert!(ys.insert(11));
1351         assert!(xs.is_disjoint(&ys));
1352         assert!(ys.is_disjoint(&xs));
1353         assert!(xs.insert(7));
1354         assert!(xs.insert(19));
1355         assert!(xs.insert(4));
1356         assert!(ys.insert(2));
1357         assert!(ys.insert(-11));
1358         assert!(xs.is_disjoint(&ys));
1359         assert!(ys.is_disjoint(&xs));
1360         assert!(ys.insert(7));
1361         assert!(!xs.is_disjoint(&ys));
1362         assert!(!ys.is_disjoint(&xs));
1363     }
1364
1365     #[test]
1366     fn test_subset_and_superset() {
1367         let mut a = TreeSet::new();
1368         assert!(a.insert(0));
1369         assert!(a.insert(5));
1370         assert!(a.insert(11));
1371         assert!(a.insert(7));
1372
1373         let mut b = TreeSet::new();
1374         assert!(b.insert(0));
1375         assert!(b.insert(7));
1376         assert!(b.insert(19));
1377         assert!(b.insert(250));
1378         assert!(b.insert(11));
1379         assert!(b.insert(200));
1380
1381         assert!(!a.is_subset(&b));
1382         assert!(!a.is_superset(&b));
1383         assert!(!b.is_subset(&a));
1384         assert!(!b.is_superset(&a));
1385
1386         assert!(b.insert(5));
1387
1388         assert!(a.is_subset(&b));
1389         assert!(!a.is_superset(&b));
1390         assert!(!b.is_subset(&a));
1391         assert!(b.is_superset(&a));
1392     }
1393
1394     #[test]
1395     fn test_iterator() {
1396         let mut m = TreeSet::new();
1397
1398         assert!(m.insert(3));
1399         assert!(m.insert(0));
1400         assert!(m.insert(4));
1401         assert!(m.insert(2));
1402         assert!(m.insert(1));
1403
1404         let mut n = 0;
1405         for x in m.iter() {
1406             assert_eq!(*x, n);
1407             n += 1
1408         }
1409     }
1410
1411     #[test]
1412     fn test_each_reverse() {
1413         let mut m = TreeSet::new();
1414
1415         assert!(m.insert(3));
1416         assert!(m.insert(0));
1417         assert!(m.insert(4));
1418         assert!(m.insert(2));
1419         assert!(m.insert(1));
1420
1421         let mut n = 4;
1422         do m.each_reverse |x| {
1423             assert_eq!(*x, n);
1424             n -= 1;
1425             true
1426         };
1427     }
1428
1429     fn check(a: &[int], b: &[int], expected: &[int],
1430              f: &fn(&TreeSet<int>, &TreeSet<int>, f: &fn(&int) -> bool) -> bool) {
1431         let mut set_a = TreeSet::new();
1432         let mut set_b = TreeSet::new();
1433
1434         for x in a.iter() { assert!(set_a.insert(*x)) }
1435         for y in b.iter() { assert!(set_b.insert(*y)) }
1436
1437         let mut i = 0;
1438         do f(&set_a, &set_b) |x| {
1439             assert_eq!(*x, expected[i]);
1440             i += 1;
1441             true
1442         };
1443         assert_eq!(i, expected.len());
1444     }
1445
1446     #[test]
1447     fn test_intersection() {
1448         fn check_intersection(a: &[int], b: &[int], expected: &[int]) {
1449             check(a, b, expected, |x, y, f| x.intersection(y).advance(f))
1450         }
1451
1452         check_intersection([], [], []);
1453         check_intersection([1, 2, 3], [], []);
1454         check_intersection([], [1, 2, 3], []);
1455         check_intersection([2], [1, 2, 3], [2]);
1456         check_intersection([1, 2, 3], [2], [2]);
1457         check_intersection([11, 1, 3, 77, 103, 5, -5],
1458                            [2, 11, 77, -9, -42, 5, 3],
1459                            [3, 5, 11, 77]);
1460     }
1461
1462     #[test]
1463     fn test_difference() {
1464         fn check_difference(a: &[int], b: &[int], expected: &[int]) {
1465             check(a, b, expected, |x, y, f| x.difference(y).advance(f))
1466         }
1467
1468         check_difference([], [], []);
1469         check_difference([1, 12], [], [1, 12]);
1470         check_difference([], [1, 2, 3, 9], []);
1471         check_difference([1, 3, 5, 9, 11],
1472                          [3, 9],
1473                          [1, 5, 11]);
1474         check_difference([-5, 11, 22, 33, 40, 42],
1475                          [-12, -5, 14, 23, 34, 38, 39, 50],
1476                          [11, 22, 33, 40, 42]);
1477     }
1478
1479     #[test]
1480     fn test_symmetric_difference() {
1481         fn check_symmetric_difference(a: &[int], b: &[int],
1482                                       expected: &[int]) {
1483             check(a, b, expected, |x, y, f| x.symmetric_difference(y).advance(f))
1484         }
1485
1486         check_symmetric_difference([], [], []);
1487         check_symmetric_difference([1, 2, 3], [2], [1, 3]);
1488         check_symmetric_difference([2], [1, 2, 3], [1, 3]);
1489         check_symmetric_difference([1, 3, 5, 9, 11],
1490                                    [-2, 3, 9, 14, 22],
1491                                    [-2, 1, 5, 11, 14, 22]);
1492     }
1493
1494     #[test]
1495     fn test_union() {
1496         fn check_union(a: &[int], b: &[int],
1497                                       expected: &[int]) {
1498             check(a, b, expected, |x, y, f| x.union(y).advance(f))
1499         }
1500
1501         check_union([], [], []);
1502         check_union([1, 2, 3], [2], [1, 2, 3]);
1503         check_union([2], [1, 2, 3], [1, 2, 3]);
1504         check_union([1, 3, 5, 9, 11, 16, 19, 24],
1505                     [-2, 1, 5, 9, 13, 19],
1506                     [-2, 1, 3, 5, 9, 11, 13, 16, 19, 24]);
1507     }
1508
1509     #[test]
1510     fn test_zip() {
1511         let mut x = TreeSet::new();
1512         x.insert(5u);
1513         x.insert(12u);
1514         x.insert(11u);
1515
1516         let mut y = TreeSet::new();
1517         y.insert("foo");
1518         y.insert("bar");
1519
1520         let x = x;
1521         let y = y;
1522         let mut z = x.iter().zip(y.iter());
1523
1524         // FIXME: #5801: this needs a type hint to compile...
1525         let result: Option<(&uint, & &'static str)> = z.next();
1526         assert_eq!(result.unwrap(), (&5u, & &"bar"));
1527
1528         let result: Option<(&uint, & &'static str)> = z.next();
1529         assert_eq!(result.unwrap(), (&11u, & &"foo"));
1530
1531         let result: Option<(&uint, & &'static str)> = z.next();
1532         assert!(result.is_none());
1533     }
1534
1535     #[test]
1536     fn test_swap() {
1537         let mut m = TreeMap::new();
1538         assert_eq!(m.swap(1, 2), None);
1539         assert_eq!(m.swap(1, 3), Some(2));
1540         assert_eq!(m.swap(1, 4), Some(3));
1541     }
1542
1543     #[test]
1544     fn test_pop() {
1545         let mut m = TreeMap::new();
1546         m.insert(1, 2);
1547         assert_eq!(m.pop(&1), Some(2));
1548         assert_eq!(m.pop(&1), None);
1549     }
1550
1551     #[test]
1552     fn test_from_iter() {
1553         let xs = ~[1, 2, 3, 4, 5, 6, 7, 8, 9];
1554
1555         let set: TreeSet<int> = xs.iter().transform(|&x| x).collect();
1556
1557         for x in xs.iter() {
1558             assert!(set.contains(x));
1559         }
1560     }
1561 }