src/libstd/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 use core::prelude::*;
  16
  17 // This is implemented as an AA tree, which is a simplified variation of
  18 // a red-black tree where where red (horizontal) nodes can only be added
  19 // as a right child. The time complexity is the same, and re-balancing
  20 // operations are more frequent but also cheaper.
  21
  22 // Future improvements:
  23
  24 // range search - O(log n) retrieval of an iterator from some key
  25
  26 // (possibly) implement the overloads Python does for sets:
  27 //   * intersection: &
  28 //   * difference: -
  29 //   * symmetric difference: ^
  30 //   * union: |
  31 // These would be convenient since the methods work like `each`
  32
  33 pub struct TreeMap<K, V> {
  34     priv root: Option<~TreeNode<K, V>>,
  35     priv length: uint
  36 }
  37
  38 impl<K: Eq + TotalOrd, V: Eq> Eq for TreeMap<K, V> {
  39     fn eq(&self, other: &TreeMap<K, V>) -> bool {
  40         if self.len() != other.len() {
  41             false
  42         } else {
  43             let mut x = self.iter();
  44             let mut y = other.iter();
  45             for self.len().times {
  46                 if map_next(&mut x).unwrap() !=
  47                    map_next(&mut y).unwrap() {
  48                     return false
  49                 }
  50             }
  51             true
  52         }
  53     }
  54     fn ne(&self, other: &TreeMap<K, V>) -> bool { !self.eq(other) }
  55 }
  56
  57 // Lexicographical comparison
  58 fn lt<K: Ord + TotalOrd, V>(a: &TreeMap<K, V>,
  59                                  b: &TreeMap<K, V>) -> bool {
  60     let mut x = a.iter();
  61     let mut y = b.iter();
  62
  63     let (a_len, b_len) = (a.len(), b.len());
  64     for uint::min(a_len, b_len).times {
  65         let (key_a,_) = map_next(&mut x).unwrap();
  66         let (key_b,_) = map_next(&mut y).unwrap();
  67         if *key_a < *key_b { return true; }
  68         if *key_a > *key_b { return false; }
  69     };
  70
  71     a_len < b_len
  72 }
  73
  74 impl<K: Ord + TotalOrd, V> Ord for TreeMap<K, V> {
  75     #[inline(always)]
  76     fn lt(&self, other: &TreeMap<K, V>) -> bool { lt(self, other) }
  77     #[inline(always)]
  78     fn le(&self, other: &TreeMap<K, V>) -> bool { !lt(other, self) }
  79     #[inline(always)]
  80     fn ge(&self, other: &TreeMap<K, V>) -> bool { !lt(self, other) }
  81     #[inline(always)]
  82     fn gt(&self, other: &TreeMap<K, V>) -> bool { lt(other, self) }
  83 }
  84
  85 impl<'self, K: TotalOrd, V> BaseIter<(&'self K, &'self V)> for TreeMap<K, V> {
  86     /// Visit all key-value pairs in order
  87     fn each(&self, f: &fn(&(&'self K, &'self V)) -> bool) {
  88         each(&self.root, f)
  89     }
  90     fn size_hint(&self) -> Option<uint> { Some(self.len()) }
  91 }
  92
  93 impl<'self, K: TotalOrd, V>
  94     ReverseIter<(&'self K, &'self V)>
  95     for TreeMap<K, V>
  96 {
  97     /// Visit all key-value pairs in reverse order
  98     fn each_reverse(&self, f: &fn(&(&'self K, &'self V)) -> bool) {
  99         each_reverse(&self.root, f);
 100     }
 101 }
 102
 103 impl<K: TotalOrd, V> Container for TreeMap<K, V> {
 104     /// Return the number of elements in the map
 105     fn len(&const self) -> uint { self.length }
 106
 107     /// Return true if the map contains no elements
 108     fn is_empty(&const self) -> bool { self.root.is_none() }
 109 }
 110
 111 impl<K: TotalOrd, V> Mutable for TreeMap<K, V> {
 112     /// Clear the map, removing all key-value pairs.
 113     fn clear(&mut self) {
 114         self.root = None;
 115         self.length = 0
 116     }
 117 }
 118
 119 impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
 120     /// Return true if the map contains a value for the specified key
 121     fn contains_key(&self, key: &K) -> bool {
 122         self.find(key).is_some()
 123     }
 124
 125     /// Visit all keys in order
 126     fn each_key(&self, f: &fn(&K) -> bool) { self.each(|&(k, _)| f(k)) }
 127
 128     /// Visit all values in order
 129     fn each_value(&self, f: &fn(&V) -> bool) {
 130         self.each(|&(_, v)| f(v))
 131     }
 132
 133     /// Iterate over the map and mutate the contained values
 134     fn mutate_values(&mut self, f: &fn(&'self K, &'self mut V) -> bool) {
 135         mutate_values(&mut self.root, f);
 136     }
 137
 138     /// Return a reference to the value corresponding to the key
 139     fn find(&self, key: &K) -> Option<&'self V> {
 140         let mut current: &'self Option<~TreeNode<K, V>> = &self.root;
 141         loop {
 142             match *current {
 143               Some(ref r) => {
 144                 match key.cmp(&r.key) {
 145                   Less => current = &r.left,
 146                   Greater => current = &r.right,
 147                   Equal => return Some(&r.value)
 148                 }
 149               }
 150               None => return None
 151             }
 152         }
 153     }
 154
 155     /// Return a mutable reference to the value corresponding to the key
 156     #[inline(always)]
 157     fn find_mut(&mut self, key: &K) -> Option<&'self mut V> {
 158         find_mut(&mut self.root, key)
 159     }
 160
 161     /// Insert a key-value pair into the map. An existing value for a
 162     /// key is replaced by the new value. Return true if the key did
 163     /// not already exist in the map.
 164     fn insert(&mut self, key: K, value: V) -> bool {
 165         let ret = insert(&mut self.root, key, value);
 166         if ret { self.length += 1 }
 167         ret
 168     }
 169
 170     /// Remove a key-value pair from the map. Return true if the key
 171     /// was present in the map, otherwise false.
 172     fn remove(&mut self, key: &K) -> bool {
 173         let ret = remove(&mut self.root, key);
 174         if ret { self.length -= 1 }
 175         ret
 176     }
 177 }
 178
 179 pub impl<K: TotalOrd, V> TreeMap<K, V> {
 180     /// Create an empty TreeMap
 181     fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} }
 182
 183     /// Visit all keys in reverse order
 184     fn each_key_reverse(&self, f: &fn(&K) -> bool) {
 185         self.each_reverse(|&(k, _)| f(k))
 186     }
 187
 188     /// Visit all values in reverse order
 189     fn each_value_reverse(&self, f: &fn(&V) -> bool) {
 190         self.each_reverse(|&(_, v)| f(v))
 191     }
 192
 193     /// Get a lazy iterator over the key-value pairs in the map.
 194     /// Requires that it be frozen (immutable).
 195     fn iter(&self) -> TreeMapIterator<'self, K, V> {
 196         TreeMapIterator{stack: ~[], node: &self.root}
 197     }
 198 }
 199
 200 /// Lazy forward iterator over a map
 201 pub struct TreeMapIterator<'self, K, V> {
 202     priv stack: ~[&'self ~TreeNode<K, V>],
 203     priv node: &'self Option<~TreeNode<K, V>>
 204 }
 205
 206 /// Advance the iterator to the next node (in order) and return a
 207 /// tuple with a reference to the key and value. If there are no
 208 /// more nodes, return `None`.
 209 pub fn map_next<'r, K, V>(iter: &mut TreeMapIterator<'r, K, V>)
 210                        -> Option<(&'r K, &'r V)> {
 211     while !iter.stack.is_empty() || iter.node.is_some() {
 212         match *iter.node {
 213           Some(ref x) => {
 214             iter.stack.push(x);
 215             iter.node = &x.left;
 216           }
 217           None => {
 218             let res = iter.stack.pop();
 219             iter.node = &res.right;
 220             return Some((&res.key, &res.value));
 221           }
 222         }
 223     }
 224     None
 225 }
 226
 227 /// Advance the iterator through the map
 228 pub fn map_advance<'r, K, V>(iter: &mut TreeMapIterator<'r, K, V>,
 229                              f: &fn((&'r K, &'r V)) -> bool) {
 230     loop {
 231         match map_next(iter) {
 232           Some(x) => {
 233             if !f(x) { return }
 234           }
 235           None => return
 236         }
 237     }
 238 }
 239
 240 pub struct TreeSet<T> {
 241     priv map: TreeMap<T, ()>
 242 }
 243
 244 impl<T: TotalOrd> BaseIter<T> for TreeSet<T> {
 245     /// Visit all values in order
 246     #[inline(always)]
 247     fn each(&self, f: &fn(&T) -> bool) { self.map.each_key(f) }
 248     #[inline(always)]
 249     fn size_hint(&self) -> Option<uint> { Some(self.len()) }
 250 }
 251
 252 impl<T: TotalOrd> ReverseIter<T> for TreeSet<T> {
 253     /// Visit all values in reverse order
 254     #[inline(always)]
 255     fn each_reverse(&self, f: &fn(&T) -> bool) {
 256         self.map.each_key_reverse(f)
 257     }
 258 }
 259
 260 impl<T: Eq + TotalOrd> Eq for TreeSet<T> {
 261     #[inline(always)]
 262     fn eq(&self, other: &TreeSet<T>) -> bool { self.map == other.map }
 263     #[inline(always)]
 264     fn ne(&self, other: &TreeSet<T>) -> bool { self.map != other.map }
 265 }
 266
 267 impl<T: Ord + TotalOrd> Ord for TreeSet<T> {
 268     #[inline(always)]
 269     fn lt(&self, other: &TreeSet<T>) -> bool { self.map < other.map }
 270     #[inline(always)]
 271     fn le(&self, other: &TreeSet<T>) -> bool { self.map <= other.map }
 272     #[inline(always)]
 273     fn ge(&self, other: &TreeSet<T>) -> bool { self.map >= other.map }
 274     #[inline(always)]
 275     fn gt(&self, other: &TreeSet<T>) -> bool { self.map > other.map }
 276 }
 277
 278 impl<T: TotalOrd> Container for TreeSet<T> {
 279     /// Return the number of elements in the set
 280     #[inline(always)]
 281     fn len(&const self) -> uint { self.map.len() }
 282
 283     /// Return true if the set contains no elements
 284     #[inline(always)]
 285     fn is_empty(&const self) -> bool { self.map.is_empty() }
 286 }
 287
 288 impl<T: TotalOrd> Mutable for TreeSet<T> {
 289     /// Clear the set, removing all values.
 290     #[inline(always)]
 291     fn clear(&mut self) { self.map.clear() }
 292 }
 293
 294 impl<T: TotalOrd> Set<T> for TreeSet<T> {
 295     /// Return true if the set contains a value
 296     #[inline(always)]
 297     fn contains(&self, value: &T) -> bool {
 298         self.map.contains_key(value)
 299     }
 300
 301     /// Add a value to the set. Return true if the value was not already
 302     /// present in the set.
 303     #[inline(always)]
 304     fn insert(&mut self, value: T) -> bool { self.map.insert(value, ()) }
 305
 306     /// Remove a value from the set. Return true if the value was
 307     /// present in the set.
 308     #[inline(always)]
 309     fn remove(&mut self, value: &T) -> bool { self.map.remove(value) }
 310
 311     /// Return true if the set has no elements in common with `other`.
 312     /// This is equivalent to checking for an empty intersection.
 313     fn is_disjoint(&self, other: &TreeSet<T>) -> bool {
 314         let mut x = self.iter();
 315         let mut y = other.iter();
 316         let mut a = set_next(&mut x);
 317         let mut b = set_next(&mut y);
 318         while a.is_some() && b.is_some() {
 319             let a1 = a.unwrap();
 320             let b1 = b.unwrap();
 321             match a1.cmp(b1) {
 322               Less => a = set_next(&mut x),
 323               Greater => b = set_next(&mut y),
 324               Equal => return false
 325             }
 326         }
 327         true
 328     }
 329
 330     /// Return true if the set is a subset of another
 331     #[inline(always)]
 332     fn is_subset(&self, other: &TreeSet<T>) -> bool {
 333         other.is_superset(self)
 334     }
 335
 336     /// Return true if the set is a superset of another
 337     fn is_superset(&self, other: &TreeSet<T>) -> bool {
 338         let mut x = self.iter();
 339         let mut y = other.iter();
 340         let mut a = set_next(&mut x);
 341         let mut b = set_next(&mut y);
 342         while b.is_some() {
 343             if a.is_none() {
 344                 return false
 345             }
 346
 347             let a1 = a.unwrap();
 348             let b1 = b.unwrap();
 349
 350             match a1.cmp(b1) {
 351               Less => (),
 352               Greater => return false,
 353               Equal => b = set_next(&mut y),
 354             }
 355
 356             a = set_next(&mut x);
 357         }
 358         true
 359     }
 360
 361     /// Visit the values (in-order) representing the difference
 362     fn difference(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) {
 363         let mut x = self.iter();
 364         let mut y = other.iter();
 365
 366         let mut a = set_next(&mut x);
 367         let mut b = set_next(&mut y);
 368
 369         while a.is_some() {
 370             if b.is_none() {
 371                 return do a.while_some() |a1| {
 372                     if f(a1) { set_next(&mut x) } else { None }
 373                 }
 374             }
 375
 376             let a1 = a.unwrap();
 377             let b1 = b.unwrap();
 378
 379             let cmp = a1.cmp(b1);
 380
 381             if cmp == Less {
 382                 if !f(a1) { return }
 383                 a = set_next(&mut x);
 384             } else {
 385                 if cmp == Equal { a = set_next(&mut x) }
 386                 b = set_next(&mut y);
 387             }
 388         }
 389     }
 390
 391     /// Visit the values (in-order) representing the symmetric difference
 392     fn symmetric_difference(&self, other: &TreeSet<T>,
 393                                  f: &fn(&T) -> bool) {
 394         let mut x = self.iter();
 395         let mut y = other.iter();
 396
 397         let mut a = set_next(&mut x);
 398         let mut b = set_next(&mut y);
 399
 400         while a.is_some() {
 401             if b.is_none() {
 402                 return do a.while_some() |a1| {
 403                     if f(a1) { set_next(&mut x) } else { None }
 404                 }
 405             }
 406
 407             let a1 = a.unwrap();
 408             let b1 = b.unwrap();
 409
 410             let cmp = a1.cmp(b1);
 411
 412             if cmp == Less {
 413                 if !f(a1) { return }
 414                 a = set_next(&mut x);
 415             } else {
 416                 if cmp == Greater {
 417                     if !f(b1) { return }
 418                 } else {
 419                     a = set_next(&mut x);
 420                 }
 421                 b = set_next(&mut y);
 422             }
 423         }
 424         do b.while_some |b1| {
 425             if f(b1) { set_next(&mut y) } else { None }
 426         }
 427     }
 428
 429     /// Visit the values (in-order) representing the intersection
 430     fn intersection(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) {
 431         let mut x = self.iter();
 432         let mut y = other.iter();
 433
 434         let mut a = set_next(&mut x);
 435         let mut b = set_next(&mut y);
 436
 437         while a.is_some() && b.is_some() {
 438             let a1 = a.unwrap();
 439             let b1 = b.unwrap();
 440
 441             let cmp = a1.cmp(b1);
 442
 443             if cmp == Less {
 444                 a = set_next(&mut x);
 445             } else {
 446                 if cmp == Equal {
 447                     if !f(a1) { return }
 448                 }
 449                 b = set_next(&mut y);
 450             }
 451         }
 452     }
 453
 454     /// Visit the values (in-order) representing the union
 455     fn union(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) {
 456         let mut x = self.iter();
 457         let mut y = other.iter();
 458
 459         let mut a = set_next(&mut x);
 460         let mut b = set_next(&mut y);
 461
 462         while a.is_some() {
 463             if b.is_none() {
 464                 return do a.while_some() |a1| {
 465                     if f(a1) { set_next(&mut x) } else { None }
 466                 }
 467             }
 468
 469             let a1 = a.unwrap();
 470             let b1 = b.unwrap();
 471
 472             let cmp = a1.cmp(b1);
 473
 474             if cmp == Greater {
 475                 if !f(b1) { return }
 476                 b = set_next(&mut y);
 477             } else {
 478                 if !f(a1) { return }
 479                 if cmp == Equal {
 480                     b = set_next(&mut y);
 481                 }
 482                 a = set_next(&mut x);
 483             }
 484         }
 485         do b.while_some |b1| {
 486             if f(b1) { set_next(&mut y) } else { None }
 487         }
 488     }
 489 }
 490
 491 pub impl <T: TotalOrd> TreeSet<T> {
 492     /// Create an empty TreeSet
 493     #[inline(always)]
 494     fn new() -> TreeSet<T> { TreeSet{map: TreeMap::new()} }
 495
 496     /// Get a lazy iterator over the values in the set.
 497     /// Requires that it be frozen (immutable).
 498     #[inline(always)]
 499     fn iter(&self) -> TreeSetIterator<'self, T> {
 500         TreeSetIterator{iter: self.map.iter()}
 501     }
 502 }
 503
 504 /// Lazy forward iterator over a set
 505 pub struct TreeSetIterator<'self, T> {
 506     priv iter: TreeMapIterator<'self, T, ()>
 507 }
 508
 509 /// Advance the iterator to the next node (in order). If this iterator is
 510 /// finished, does nothing.
 511 #[inline(always)]
 512 pub fn set_next<'r, T>(iter: &mut TreeSetIterator<'r, T>) -> Option<&'r T> {
 513     do map_next(&mut iter.iter).map |&(value, _)| { value }
 514 }
 515
 516 /// Advance the iterator through the set
 517 #[inline(always)]
 518 pub fn set_advance<'r, T>(iter: &mut TreeSetIterator<'r, T>,
 519                           f: &fn(&'r T) -> bool) {
 520     do map_advance(&mut iter.iter) |(k, _)| { f(k) }
 521 }
 522
 523 // Nodes keep track of their level in the tree, starting at 1 in the
 524 // leaves and with a red child sharing the level of the parent.
 525 struct TreeNode<K, V> {
 526     key: K,
 527     value: V,
 528     left: Option<~TreeNode<K, V>>,
 529     right: Option<~TreeNode<K, V>>,
 530     level: uint
 531 }
 532
 533 pub impl<K: TotalOrd, V> TreeNode<K, V> {
 534     #[inline(always)]
 535     fn new(key: K, value: V) -> TreeNode<K, V> {
 536         TreeNode{key: key, value: value, left: None, right: None, level: 1}
 537     }
 538 }
 539
 540 fn each<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>,
 541                             f: &fn(&(&'r K, &'r V)) -> bool) {
 542     for node.each |x| {
 543         each(&x.left, f);
 544         if f(&(&x.key, &x.value)) { each(&x.right, f) }
 545     }
 546 }
 547
 548 fn each_reverse<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>,
 549                                     f: &fn(&(&'r K, &'r V)) -> bool) {
 550     for node.each |x| {
 551         each_reverse(&x.right, f);
 552         if f(&(&x.key, &x.value)) { each_reverse(&x.left, f) }
 553     }
 554 }
 555
 556 fn mutate_values<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
 557                                      f: &fn(&'r K, &'r mut V) -> bool)
 558                                   -> bool {
 559     match *node {
 560       Some(~TreeNode{key: ref key, value: ref mut value, left: ref mut left,
 561                      right: ref mut right, _}) => {
 562         if !mutate_values(left, f) { return false }
 563         if !f(key, value) { return false }
 564         if !mutate_values(right, f) { return false }
 565       }
 566       None => return false
 567     }
 568     true
 569 }
 570
 571 // Remove left horizontal link by rotating right
 572 fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
 573     if node.left.map_default(false, |x| x.level == node.level) {
 574         let mut save = node.left.swap_unwrap();
 575         node.left <-> save.right; // save.right now None
 576         *node <-> save;
 577         node.right = Some(save);
 578     }
 579 }
 580
 581 // Remove dual horizontal link by rotating left and increasing level of
 582 // the parent
 583 fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
 584     if node.right.map_default(false,
 585       |x| x.right.map_default(false, |y| y.level == node.level)) {
 586         let mut save = node.right.swap_unwrap();
 587         node.right <-> save.left; // save.left now None
 588         save.level += 1;
 589         *node <-> save;
 590         node.left = Some(save);
 591     }
 592 }
 593
 594 fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
 595                                 key: &K)
 596                              -> Option<&'r mut V> {
 597     match *node {
 598       Some(ref mut x) => {
 599         match key.cmp(&x.key) {
 600           Less => find_mut(&mut x.left, key),
 601           Greater => find_mut(&mut x.right, key),
 602           Equal => Some(&mut x.value),
 603         }
 604       }
 605       None => None
 606     }
 607 }
 608
 609 fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K, value: V) -> bool {
 610     match *node {
 611       Some(ref mut save) => {
 612         match key.cmp(&save.key) {
 613           Less => {
 614             let inserted = insert(&mut save.left, key, value);
 615             skew(save);
 616             split(save);
 617             inserted
 618           }
 619           Greater => {
 620             let inserted = insert(&mut save.right, key, value);
 621             skew(save);
 622             split(save);
 623             inserted
 624           }
 625           Equal => {
 626             save.key = key;
 627             save.value = value;
 628             false
 629           }
 630         }
 631       }
 632       None => {
 633        *node = Some(~TreeNode::new(key, value));
 634         true
 635       }
 636     }
 637 }
 638
 639 fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
 640                           key: &K) -> bool {
 641     fn heir_swap<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>,
 642                             child: &mut Option<~TreeNode<K, V>>) {
 643         // *could* be done without recursion, but it won't borrow check
 644         for child.each_mut |x| {
 645             if x.right.is_some() {
 646                 heir_swap(node, &mut x.right);
 647             } else {
 648                 node.key <-> x.key;
 649                 node.value <-> x.value;
 650             }
 651         }
 652     }
 653
 654     match *node {
 655       None => {
 656         return false // bottom of tree
 657       }
 658       Some(ref mut save) => {
 659         let (removed, this) = match key.cmp(&save.key) {
 660           Less => (remove(&mut save.left, key), false),
 661           Greater => (remove(&mut save.right, key), false),
 662           Equal => {
 663             if save.left.is_some() {
 664                 if save.right.is_some() {
 665                     let mut left = save.left.swap_unwrap();
 666                     if left.right.is_some() {
 667                         heir_swap(save, &mut left.right);
 668                     } else {
 669                         save.key <-> left.key;
 670                         save.value <-> left.value;
 671                     }
 672                     save.left = Some(left);
 673                     remove(&mut save.left, key);
 674                 } else {
 675                     *save = save.left.swap_unwrap();
 676                 }
 677                 (true, false)
 678             } else if save.right.is_some() {
 679                 *save = save.right.swap_unwrap();
 680                 (true, false)
 681             } else {
 682                 (true, true)
 683             }
 684           }
 685         };
 686
 687         if !this {
 688             let left_level = save.left.map_default(0, |x| x.level);
 689             let right_level = save.right.map_default(0, |x| x.level);
 690
 691             // re-balance, if necessary
 692             if left_level < save.level - 1 || right_level < save.level - 1 {
 693                 save.level -= 1;
 694
 695                 if right_level > save.level {
 696                     for save.right.each_mut |x| { x.level = save.level }
 697                 }
 698
 699                 skew(save);
 700
 701                 for save.right.each_mut |right| {
 702                     skew(right);
 703                     for right.right.each_mut |x| { skew(x) }
 704                 }
 705
 706                 split(save);
 707                 for save.right.each_mut |x| { split(x) }
 708             }
 709
 710             return removed;
 711         }
 712       }
 713     }
 714
 715     *node = None;
 716     true
 717 }
 718
 719 #[cfg(test)]
 720 mod test_treemap {
 721     use core::prelude::*;
 722     use super::*;
 723     use core::rand::RngUtil;
 724     use core::rand;
 725
 726     #[test]
 727     fn find_empty() {
 728         let m = TreeMap::new::<int, int>(); assert!(m.find(&5) == None);
 729     }
 730
 731     #[test]
 732     fn find_not_found() {
 733         let mut m = TreeMap::new();
 734         assert!(m.insert(1, 2));
 735         assert!(m.insert(5, 3));
 736         assert!(m.insert(9, 3));
 737         assert!(m.find(&2) == None);
 738     }
 739
 740     #[test]
 741     fn test_find_mut() {
 742         let mut m = TreeMap::new();
 743         assert!(m.insert(1, 12));
 744         assert!(m.insert(2, 8));
 745         assert!(m.insert(5, 14));
 746         let new = 100;
 747         match m.find_mut(&5) {
 748           None => fail!(), Some(x) => *x = new
 749         }
 750         assert_eq!(m.find(&5), Some(&new));
 751     }
 752
 753     #[test]
 754     fn insert_replace() {
 755         let mut m = TreeMap::new();
 756         assert!(m.insert(5, 2));
 757         assert!(m.insert(2, 9));
 758         assert!(!m.insert(2, 11));
 759         assert!(m.find(&2).unwrap() == &11);
 760     }
 761
 762     #[test]
 763     fn test_clear() {
 764         let mut m = TreeMap::new();
 765         m.clear();
 766         assert!(m.insert(5, 11));
 767         assert!(m.insert(12, -3));
 768         assert!(m.insert(19, 2));
 769         m.clear();
 770         assert!(m.find(&5).is_none());
 771         assert!(m.find(&12).is_none());
 772         assert!(m.find(&19).is_none());
 773         assert!(m.is_empty());
 774     }
 775
 776     #[test]
 777     fn u8_map() {
 778         let mut m = TreeMap::new();
 779
 780         let k1 = str::to_bytes(~"foo");
 781         let k2 = str::to_bytes(~"bar");
 782         let v1 = str::to_bytes(~"baz");
 783         let v2 = str::to_bytes(~"foobar");
 784
 785         m.insert(copy k1, copy v1);
 786         m.insert(copy k2, copy v2);
 787
 788         assert!(m.find(&k2) == Some(&v2));
 789         assert!(m.find(&k1) == Some(&v1));
 790     }
 791
 792     fn check_equal<K: Eq + TotalOrd, V: Eq>(ctrl: &[(K, V)],
 793                                             map: &TreeMap<K, V>) {
 794         assert!(ctrl.is_empty() == map.is_empty());
 795         for ctrl.each |x| {
 796             let &(k, v) = x;
 797             assert!(map.find(&k).unwrap() == &v)
 798         }
 799         for map.each |&(map_k, map_v)| {
 800             let mut found = false;
 801             for ctrl.each |x| {
 802                 let &(ctrl_k, ctrl_v) = x;
 803                 if *map_k == ctrl_k {
 804                     assert!(*map_v == ctrl_v);
 805                     found = true;
 806                     break;
 807                 }
 808             }
 809             assert!(found);
 810         }
 811     }
 812
 813     fn check_left<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>,
 814                                   parent: &~TreeNode<K, V>) {
 815         match *node {
 816           Some(ref r) => {
 817             assert!(r.key.cmp(&parent.key) == Less);
 818             assert!(r.level == parent.level - 1); // left is black
 819             check_left(&r.left, r);
 820             check_right(&r.right, r, false);
 821           }
 822           None => assert!(parent.level == 1) // parent is leaf
 823         }
 824     }
 825
 826     fn check_right<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>,
 827                                    parent: &~TreeNode<K, V>,
 828                                    parent_red: bool) {
 829         match *node {
 830           Some(ref r) => {
 831             assert!(r.key.cmp(&parent.key) == Greater);
 832             let red = r.level == parent.level;
 833             if parent_red { assert!(!red) } // no dual horizontal links
 834             // Right red or black
 835             assert!(red || r.level == parent.level - 1);
 836             check_left(&r.left, r);
 837             check_right(&r.right, r, red);
 838           }
 839           None => assert!(parent.level == 1) // parent is leaf
 840         }
 841     }
 842
 843     fn check_structure<K: TotalOrd, V>(map: &TreeMap<K, V>) {
 844         match map.root {
 845           Some(ref r) => {
 846             check_left(&r.left, r);
 847             check_right(&r.right, r, false);
 848           }
 849           None => ()
 850         }
 851     }
 852
 853     #[test]
 854     fn test_rand_int() {
 855         let mut map = TreeMap::new::<int, int>();
 856         let mut ctrl = ~[];
 857
 858         check_equal(ctrl, &map);
 859         assert!(map.find(&5).is_none());
 860
 861         let rng = rand::seeded_rng(&[42]);
 862
 863         for 3.times {
 864             for 90.times {
 865                 let k = rng.gen_int();
 866                 let v = rng.gen_int();
 867                 if !ctrl.contains(&(k, v)) {
 868                     assert!(map.insert(k, v));
 869                     ctrl.push((k, v));
 870                     check_structure(&map);
 871                     check_equal(ctrl, &map);
 872                 }
 873             }
 874
 875             for 30.times {
 876                 let r = rng.gen_uint_range(0, ctrl.len());
 877                 let (key, _) = vec::remove(&mut ctrl, r);
 878                 assert!(map.remove(&key));
 879                 check_structure(&map);
 880                 check_equal(ctrl, &map);
 881             }
 882         }
 883     }
 884
 885     #[test]
 886     fn test_len() {
 887         let mut m = TreeMap::new();
 888         assert!(m.insert(3, 6));
 889         assert!(m.len() == 1);
 890         assert!(m.insert(0, 0));
 891         assert!(m.len() == 2);
 892         assert!(m.insert(4, 8));
 893         assert!(m.len() == 3);
 894         assert!(m.remove(&3));
 895         assert!(m.len() == 2);
 896         assert!(!m.remove(&5));
 897         assert!(m.len() == 2);
 898         assert!(m.insert(2, 4));
 899         assert!(m.len() == 3);
 900         assert!(m.insert(1, 2));
 901         assert!(m.len() == 4);
 902     }
 903
 904     #[test]
 905     fn test_each() {
 906         let mut m = TreeMap::new();
 907
 908         assert!(m.insert(3, 6));
 909         assert!(m.insert(0, 0));
 910         assert!(m.insert(4, 8));
 911         assert!(m.insert(2, 4));
 912         assert!(m.insert(1, 2));
 913
 914         let mut n = 0;
 915         for m.each |&(k, v)| {
 916             assert!(*k == n);
 917             assert!(*v == n * 2);
 918             n += 1;
 919         }
 920     }
 921
 922     #[test]
 923     fn test_each_reverse() {
 924         let mut m = TreeMap::new();
 925
 926         assert!(m.insert(3, 6));
 927         assert!(m.insert(0, 0));
 928         assert!(m.insert(4, 8));
 929         assert!(m.insert(2, 4));
 930         assert!(m.insert(1, 2));
 931
 932         let mut n = 4;
 933         for m.each_reverse |&(k, v)| {
 934             assert!(*k == n);
 935             assert!(*v == n * 2);
 936             n -= 1;
 937         }
 938     }
 939
 940     #[test]
 941     fn test_eq() {
 942         let mut a = TreeMap::new();
 943         let mut b = TreeMap::new();
 944
 945         assert!(a == b);
 946         assert!(a.insert(0, 5));
 947         assert!(a != b);
 948         assert!(b.insert(0, 4));
 949         assert!(a != b);
 950         assert!(a.insert(5, 19));
 951         assert!(a != b);
 952         assert!(!b.insert(0, 5));
 953         assert!(a != b);
 954         assert!(b.insert(5, 19));
 955         assert!(a == b);
 956     }
 957
 958     #[test]
 959     fn test_lt() {
 960         let mut a = TreeMap::new();
 961         let mut b = TreeMap::new();
 962
 963         assert!(!(a < b) && !(b < a));
 964         assert!(b.insert(0, 5));
 965         assert!(a < b);
 966         assert!(a.insert(0, 7));
 967         assert!(!(a < b) && !(b < a));
 968         assert!(b.insert(-2, 0));
 969         assert!(b < a);
 970         assert!(a.insert(-5, 2));
 971         assert!(a < b);
 972         assert!(a.insert(6, 2));
 973         assert!(a < b && !(b < a));
 974     }
 975
 976     #[test]
 977     fn test_ord() {
 978         let mut a = TreeMap::new();
 979         let mut b = TreeMap::new();
 980
 981         assert!(a <= b && a >= b);
 982         assert!(a.insert(1, 1));
 983         assert!(a > b && a >= b);
 984         assert!(b < a && b <= a);
 985         assert!(b.insert(2, 2));
 986         assert!(b > a && b >= a);
 987         assert!(a < b && a <= b);
 988     }
 989
 990     #[test]
 991     fn test_lazy_iterator() {
 992         let mut m = TreeMap::new();
 993         let (x1, y1) = (2, 5);
 994         let (x2, y2) = (9, 12);
 995         let (x3, y3) = (20, -3);
 996         let (x4, y4) = (29, 5);
 997         let (x5, y5) = (103, 3);
 998
 999         assert!(m.insert(x1, y1));
1000         assert!(m.insert(x2, y2));
1001         assert!(m.insert(x3, y3));
1002         assert!(m.insert(x4, y4));
1003         assert!(m.insert(x5, y5));
1004
1005         let m = m;
1006         let mut a = m.iter();
1007
1008         assert!(map_next(&mut a).unwrap() == (&x1, &y1));
1009         assert!(map_next(&mut a).unwrap() == (&x2, &y2));
1010         assert!(map_next(&mut a).unwrap() == (&x3, &y3));
1011         assert!(map_next(&mut a).unwrap() == (&x4, &y4));
1012         assert!(map_next(&mut a).unwrap() == (&x5, &y5));
1013
1014         assert!(map_next(&mut a).is_none());
1015
1016         let mut b = m.iter();
1017
1018         let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4),
1019                         (&x5, &y5)];
1020         let mut i = 0;
1021
1022         for map_advance(&mut b) |x| {
1023             assert!(expected[i] == x);
1024             i += 1;
1025
1026             if i == 2 {
1027                 break
1028             }
1029         }
1030
1031         for map_advance(&mut b) |x| {
1032             assert!(expected[i] == x);
1033             i += 1;
1034         }
1035     }
1036 }
1037
1038 #[cfg(test)]
1039 mod test_set {
1040     use super::*;
1041
1042     #[test]
1043     fn test_clear() {
1044         let mut s = TreeSet::new();
1045         s.clear();
1046         assert!(s.insert(5));
1047         assert!(s.insert(12));
1048         assert!(s.insert(19));
1049         s.clear();
1050         assert!(!s.contains(&5));
1051         assert!(!s.contains(&12));
1052         assert!(!s.contains(&19));
1053         assert!(s.is_empty());
1054     }
1055
1056     #[test]
1057     fn test_disjoint() {
1058         let mut xs = TreeSet::new();
1059         let mut ys = TreeSet::new();
1060         assert!(xs.is_disjoint(&ys));
1061         assert!(ys.is_disjoint(&xs));
1062         assert!(xs.insert(5));
1063         assert!(ys.insert(11));
1064         assert!(xs.is_disjoint(&ys));
1065         assert!(ys.is_disjoint(&xs));
1066         assert!(xs.insert(7));
1067         assert!(xs.insert(19));
1068         assert!(xs.insert(4));
1069         assert!(ys.insert(2));
1070         assert!(ys.insert(-11));
1071         assert!(xs.is_disjoint(&ys));
1072         assert!(ys.is_disjoint(&xs));
1073         assert!(ys.insert(7));
1074         assert!(!xs.is_disjoint(&ys));
1075         assert!(!ys.is_disjoint(&xs));
1076     }
1077
1078     #[test]
1079     fn test_subset_and_superset() {
1080         let mut a = TreeSet::new();
1081         assert!(a.insert(0));
1082         assert!(a.insert(5));
1083         assert!(a.insert(11));
1084         assert!(a.insert(7));
1085
1086         let mut b = TreeSet::new();
1087         assert!(b.insert(0));
1088         assert!(b.insert(7));
1089         assert!(b.insert(19));
1090         assert!(b.insert(250));
1091         assert!(b.insert(11));
1092         assert!(b.insert(200));
1093
1094         assert!(!a.is_subset(&b));
1095         assert!(!a.is_superset(&b));
1096         assert!(!b.is_subset(&a));
1097         assert!(!b.is_superset(&a));
1098
1099         assert!(b.insert(5));
1100
1101         assert!(a.is_subset(&b));
1102         assert!(!a.is_superset(&b));
1103         assert!(!b.is_subset(&a));
1104         assert!(b.is_superset(&a));
1105     }
1106
1107     #[test]
1108     fn test_each() {
1109         let mut m = TreeSet::new();
1110
1111         assert!(m.insert(3));
1112         assert!(m.insert(0));
1113         assert!(m.insert(4));
1114         assert!(m.insert(2));
1115         assert!(m.insert(1));
1116
1117         let mut n = 0;
1118         for m.each |x| {
1119             assert!(*x == n);
1120             n += 1
1121         }
1122     }
1123
1124     #[test]
1125     fn test_each_reverse() {
1126         let mut m = TreeSet::new();
1127
1128         assert!(m.insert(3));
1129         assert!(m.insert(0));
1130         assert!(m.insert(4));
1131         assert!(m.insert(2));
1132         assert!(m.insert(1));
1133
1134         let mut n = 4;
1135         for m.each_reverse |x| {
1136             assert!(*x == n);
1137             n -= 1
1138         }
1139     }
1140
1141     fn check(a: &[int], b: &[int], expected: &[int],
1142              f: &fn(&TreeSet<int>, &TreeSet<int>, f: &fn(&int) -> bool)) {
1143         let mut set_a = TreeSet::new();
1144         let mut set_b = TreeSet::new();
1145
1146         for a.each |x| { assert!(set_a.insert(*x)) }
1147         for b.each |y| { assert!(set_b.insert(*y)) }
1148
1149         let mut i = 0;
1150         for f(&set_a, &set_b) |x| {
1151             assert!(*x == expected[i]);
1152             i += 1;
1153         }
1154         assert!(i == expected.len());
1155     }
1156
1157     #[test]
1158     fn test_intersection() {
1159         fn check_intersection(a: &[int], b: &[int], expected: &[int]) {
1160             check(a, b, expected, |x, y, z| x.intersection(y, z))
1161         }
1162
1163         check_intersection([], [], []);
1164         check_intersection([1, 2, 3], [], []);
1165         check_intersection([], [1, 2, 3], []);
1166         check_intersection([2], [1, 2, 3], [2]);
1167         check_intersection([1, 2, 3], [2], [2]);
1168         check_intersection([11, 1, 3, 77, 103, 5, -5],
1169                            [2, 11, 77, -9, -42, 5, 3],
1170                            [3, 5, 11, 77]);
1171     }
1172
1173     #[test]
1174     fn test_difference() {
1175         fn check_difference(a: &[int], b: &[int], expected: &[int]) {
1176             check(a, b, expected, |x, y, z| x.difference(y, z))
1177         }
1178
1179         check_difference([], [], []);
1180         check_difference([1, 12], [], [1, 12]);
1181         check_difference([], [1, 2, 3, 9], []);
1182         check_difference([1, 3, 5, 9, 11],
1183                          [3, 9],
1184                          [1, 5, 11]);
1185         check_difference([-5, 11, 22, 33, 40, 42],
1186                          [-12, -5, 14, 23, 34, 38, 39, 50],
1187                          [11, 22, 33, 40, 42]);
1188     }
1189
1190     #[test]
1191     fn test_symmetric_difference() {
1192         fn check_symmetric_difference(a: &[int], b: &[int],
1193                                       expected: &[int]) {
1194             check(a, b, expected, |x, y, z| x.symmetric_difference(y, z))
1195         }
1196
1197         check_symmetric_difference([], [], []);
1198         check_symmetric_difference([1, 2, 3], [2], [1, 3]);
1199         check_symmetric_difference([2], [1, 2, 3], [1, 3]);
1200         check_symmetric_difference([1, 3, 5, 9, 11],
1201                                    [-2, 3, 9, 14, 22],
1202                                    [-2, 1, 5, 11, 14, 22]);
1203     }
1204
1205     #[test]
1206     fn test_union() {
1207         fn check_union(a: &[int], b: &[int],
1208                                       expected: &[int]) {
1209             check(a, b, expected, |x, y, z| x.union(y, z))
1210         }
1211
1212         check_union([], [], []);
1213         check_union([1, 2, 3], [2], [1, 2, 3]);
1214         check_union([2], [1, 2, 3], [1, 2, 3]);
1215         check_union([1, 3, 5, 9, 11, 16, 19, 24],
1216                     [-2, 1, 5, 9, 13, 19],
1217                     [-2, 1, 3, 5, 9, 11, 13, 16, 19, 24]);
1218     }
1219 }