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