src/libcollections/trie.rs

   1 // Copyright 2013-2014 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 //! Ordered containers with integer keys, implemented as radix tries (`TrieSet` and `TrieMap` types)
  12
  13 use core::prelude::*;
  14
  15 use alloc::boxed::Box;
  16 use core::default::Default;
  17 use core::mem::zeroed;
  18 use core::mem;
  19 use core::uint;
  20
  21 use {Collection, Mutable, Map, MutableMap, Set, MutableSet};
  22 use slice::{Items, MutItems};
  23 use slice;
  24
  25 // FIXME: #5244: need to manually update the TrieNode constructor
  26 static SHIFT: uint = 4;
  27 static SIZE: uint = 1 << SHIFT;
  28 static MASK: uint = SIZE - 1;
  29 static NUM_CHUNKS: uint = uint::BITS / SHIFT;
  30
  31 enum Child<T> {
  32     Internal(Box<TrieNode<T>>),
  33     External(uint, T),
  34     Nothing
  35 }
  36
  37 #[allow(missing_doc)]
  38 pub struct TrieMap<T> {
  39     root: TrieNode<T>,
  40     length: uint
  41 }
  42
  43 impl<T> Collection for TrieMap<T> {
  44     /// Return the number of elements in the map
  45     #[inline]
  46     fn len(&self) -> uint { self.length }
  47 }
  48
  49 impl<T> Mutable for TrieMap<T> {
  50     /// Clear the map, removing all values.
  51     #[inline]
  52     fn clear(&mut self) {
  53         self.root = TrieNode::new();
  54         self.length = 0;
  55     }
  56 }
  57
  58 impl<T> Map<uint, T> for TrieMap<T> {
  59     /// Return a reference to the value corresponding to the key
  60     #[inline]
  61     fn find<'a>(&'a self, key: &uint) -> Option<&'a T> {
  62         let mut node: &'a TrieNode<T> = &self.root;
  63         let mut idx = 0;
  64         loop {
  65             match node.children[chunk(*key, idx)] {
  66               Internal(ref x) => node = &**x,
  67               External(stored, ref value) => {
  68                 if stored == *key {
  69                     return Some(value)
  70                 } else {
  71                     return None
  72                 }
  73               }
  74               Nothing => return None
  75             }
  76             idx += 1;
  77         }
  78     }
  79 }
  80
  81 impl<T> MutableMap<uint, T> for TrieMap<T> {
  82     /// Return a mutable reference to the value corresponding to the key
  83     #[inline]
  84     fn find_mut<'a>(&'a mut self, key: &uint) -> Option<&'a mut T> {
  85         find_mut(&mut self.root.children[chunk(*key, 0)], *key, 1)
  86     }
  87
  88     /// Insert a key-value pair from the map. If the key already had a value
  89     /// present in the map, that value is returned. Otherwise None is returned.
  90     fn swap(&mut self, key: uint, value: T) -> Option<T> {
  91         let ret = insert(&mut self.root.count,
  92                          &mut self.root.children[chunk(key, 0)],
  93                          key, value, 1);
  94         if ret.is_none() { self.length += 1 }
  95         ret
  96     }
  97
  98     /// Removes a key from the map, returning the value at the key if the key
  99     /// was previously in the map.
 100     fn pop(&mut self, key: &uint) -> Option<T> {
 101         let ret = remove(&mut self.root.count,
 102                          &mut self.root.children[chunk(*key, 0)],
 103                          *key, 1);
 104         if ret.is_some() { self.length -= 1 }
 105         ret
 106     }
 107 }
 108
 109 impl<T> Default for TrieMap<T> {
 110     #[inline]
 111     fn default() -> TrieMap<T> { TrieMap::new() }
 112 }
 113
 114 impl<T> TrieMap<T> {
 115     /// Create an empty TrieMap
 116     #[inline]
 117     pub fn new() -> TrieMap<T> {
 118         TrieMap{root: TrieNode::new(), length: 0}
 119     }
 120
 121     /// Visit all key-value pairs in reverse order
 122     #[inline]
 123     pub fn each_reverse<'a>(&'a self, f: |&uint, &'a T| -> bool) -> bool {
 124         self.root.each_reverse(f)
 125     }
 126
 127     /// Get an iterator over the key-value pairs in the map
 128     pub fn iter<'a>(&'a self) -> Entries<'a, T> {
 129         let mut iter = unsafe {Entries::new()};
 130         iter.stack[0] = self.root.children.iter();
 131         iter.length = 1;
 132         iter.remaining_min = self.length;
 133         iter.remaining_max = self.length;
 134
 135         iter
 136     }
 137
 138     /// Get an iterator over the key-value pairs in the map, with the
 139     /// ability to mutate the values.
 140     pub fn mut_iter<'a>(&'a mut self) -> MutEntries<'a, T> {
 141         let mut iter = unsafe {MutEntries::new()};
 142         iter.stack[0] = self.root.children.mut_iter();
 143         iter.length = 1;
 144         iter.remaining_min = self.length;
 145         iter.remaining_max = self.length;
 146
 147         iter
 148     }
 149 }
 150
 151 // FIXME #5846 we want to be able to choose between &x and &mut x
 152 // (with many different `x`) below, so we need to optionally pass mut
 153 // as a tt, but the only thing we can do with a `tt` is pass them to
 154 // other macros, so this takes the `& <mutability> <operand>` token
 155 // sequence and forces their evaluation as an expression. (see also
 156 // `item!` below.)
 157 macro_rules! addr { ($e:expr) => { $e } }
 158
 159 macro_rules! bound {
 160     ($iterator_name:ident,
 161      // the current treemap
 162      self = $this:expr,
 163      // the key to look for
 164      key = $key:expr,
 165      // are we looking at the upper bound?
 166      is_upper = $upper:expr,
 167
 168      // method names for slicing/iterating.
 169      slice_from = $slice_from:ident,
 170      iter = $iter:ident,
 171
 172      // see the comment on `addr!`, this is just an optional mut, but
 173      // there's no 0-or-1 repeats yet.
 174      mutability = $($mut_:tt)*) => {
 175         {
 176             // # For `mut`
 177             // We need an unsafe pointer here because we are borrowing
 178             // mutable references to the internals of each of these
 179             // mutable nodes, while still using the outer node.
 180             //
 181             // However, we're allowed to flaunt rustc like this because we
 182             // never actually modify the "shape" of the nodes. The only
 183             // place that mutation is can actually occur is of the actual
 184             // values of the TrieMap (as the return value of the
 185             // iterator), i.e. we can never cause a deallocation of any
 186             // TrieNodes so the raw pointer is always valid.
 187             //
 188             // # For non-`mut`
 189             // We like sharing code so much that even a little unsafe won't
 190             // stop us.
 191             let this = $this;
 192             let mut node = unsafe {
 193                 mem::transmute::<_, uint>(&this.root) as *mut TrieNode<T>
 194             };
 195
 196             let key = $key;
 197
 198             let mut it = unsafe {$iterator_name::new()};
 199             // everything else is zero'd, as we want.
 200             it.remaining_max = this.length;
 201
 202             // this addr is necessary for the `Internal` pattern.
 203             addr!(loop {
 204                     let children = unsafe {addr!(& $($mut_)* (*node).children)};
 205                     // it.length is the current depth in the iterator and the
 206                     // current depth through the `uint` key we've traversed.
 207                     let child_id = chunk(key, it.length);
 208                     let (slice_idx, ret) = match children[child_id] {
 209                         Internal(ref $($mut_)* n) => {
 210                             node = unsafe {
 211                                 mem::transmute::<_, uint>(&**n)
 212                                     as *mut TrieNode<T>
 213                             };
 214                             (child_id + 1, false)
 215                         }
 216                         External(stored, _) => {
 217                             (if stored < key || ($upper && stored == key) {
 218                                 child_id + 1
 219                             } else {
 220                                 child_id
 221                             }, true)
 222                         }
 223                         Nothing => {
 224                             (child_id + 1, true)
 225                         }
 226                     };
 227                     // push to the stack.
 228                     it.stack[it.length] = children.$slice_from(slice_idx).$iter();
 229                     it.length += 1;
 230                     if ret { return it }
 231                 })
 232         }
 233     }
 234 }
 235
 236 impl<T> TrieMap<T> {
 237     // If `upper` is true then returns upper_bound else returns lower_bound.
 238     #[inline]
 239     fn bound<'a>(&'a self, key: uint, upper: bool) -> Entries<'a, T> {
 240         bound!(Entries, self = self,
 241                key = key, is_upper = upper,
 242                slice_from = slice_from, iter = iter,
 243                mutability = )
 244     }
 245
 246     /// Get an iterator pointing to the first key-value pair whose key is not less than `key`.
 247     /// If all keys in the map are less than `key` an empty iterator is returned.
 248     pub fn lower_bound<'a>(&'a self, key: uint) -> Entries<'a, T> {
 249         self.bound(key, false)
 250     }
 251
 252     /// Get an iterator pointing to the first key-value pair whose key is greater than `key`.
 253     /// If all keys in the map are not greater than `key` an empty iterator is returned.
 254     pub fn upper_bound<'a>(&'a self, key: uint) -> Entries<'a, T> {
 255         self.bound(key, true)
 256     }
 257     // If `upper` is true then returns upper_bound else returns lower_bound.
 258     #[inline]
 259     fn mut_bound<'a>(&'a mut self, key: uint, upper: bool) -> MutEntries<'a, T> {
 260         bound!(MutEntries, self = self,
 261                key = key, is_upper = upper,
 262                slice_from = mut_slice_from, iter = mut_iter,
 263                mutability = mut)
 264     }
 265
 266     /// Get an iterator pointing to the first key-value pair whose key is not less than `key`.
 267     /// If all keys in the map are less than `key` an empty iterator is returned.
 268     pub fn mut_lower_bound<'a>(&'a mut self, key: uint) -> MutEntries<'a, T> {
 269         self.mut_bound(key, false)
 270     }
 271
 272     /// Get an iterator pointing to the first key-value pair whose key is greater than `key`.
 273     /// If all keys in the map are not greater than `key` an empty iterator is returned.
 274     pub fn mut_upper_bound<'a>(&'a mut self, key: uint) -> MutEntries<'a, T> {
 275         self.mut_bound(key, true)
 276     }
 277 }
 278
 279 impl<T> FromIterator<(uint, T)> for TrieMap<T> {
 280     fn from_iter<Iter: Iterator<(uint, T)>>(iter: Iter) -> TrieMap<T> {
 281         let mut map = TrieMap::new();
 282         map.extend(iter);
 283         map
 284     }
 285 }
 286
 287 impl<T> Extendable<(uint, T)> for TrieMap<T> {
 288     fn extend<Iter: Iterator<(uint, T)>>(&mut self, mut iter: Iter) {
 289         for (k, v) in iter {
 290             self.insert(k, v);
 291         }
 292     }
 293 }
 294
 295 #[allow(missing_doc)]
 296 pub struct TrieSet {
 297     map: TrieMap<()>
 298 }
 299
 300 impl Collection for TrieSet {
 301     /// Return the number of elements in the set
 302     #[inline]
 303     fn len(&self) -> uint { self.map.len() }
 304 }
 305
 306 impl Mutable for TrieSet {
 307     /// Clear the set, removing all values.
 308     #[inline]
 309     fn clear(&mut self) { self.map.clear() }
 310 }
 311
 312 impl Set<uint> for TrieSet {
 313     #[inline]
 314     fn contains(&self, value: &uint) -> bool {
 315         self.map.contains_key(value)
 316     }
 317
 318     #[inline]
 319     fn is_disjoint(&self, other: &TrieSet) -> bool {
 320         self.iter().all(|v| !other.contains(&v))
 321     }
 322
 323     #[inline]
 324     fn is_subset(&self, other: &TrieSet) -> bool {
 325         self.iter().all(|v| other.contains(&v))
 326     }
 327
 328     #[inline]
 329     fn is_superset(&self, other: &TrieSet) -> bool {
 330         other.is_subset(self)
 331     }
 332 }
 333
 334 impl MutableSet<uint> for TrieSet {
 335     #[inline]
 336     fn insert(&mut self, value: uint) -> bool {
 337         self.map.insert(value, ())
 338     }
 339
 340     #[inline]
 341     fn remove(&mut self, value: &uint) -> bool {
 342         self.map.remove(value)
 343     }
 344 }
 345
 346 impl Default for TrieSet {
 347     #[inline]
 348     fn default() -> TrieSet { TrieSet::new() }
 349 }
 350
 351 impl TrieSet {
 352     /// Create an empty TrieSet
 353     #[inline]
 354     pub fn new() -> TrieSet {
 355         TrieSet{map: TrieMap::new()}
 356     }
 357
 358     /// Visit all values in reverse order
 359     #[inline]
 360     pub fn each_reverse(&self, f: |&uint| -> bool) -> bool {
 361         self.map.each_reverse(|k, _| f(k))
 362     }
 363
 364     /// Get an iterator over the values in the set
 365     #[inline]
 366     pub fn iter<'a>(&'a self) -> SetItems<'a> {
 367         SetItems{iter: self.map.iter()}
 368     }
 369
 370     /// Get an iterator pointing to the first value that is not less than `val`.
 371     /// If all values in the set are less than `val` an empty iterator is returned.
 372     pub fn lower_bound<'a>(&'a self, val: uint) -> SetItems<'a> {
 373         SetItems{iter: self.map.lower_bound(val)}
 374     }
 375
 376     /// Get an iterator pointing to the first value that key is greater than `val`.
 377     /// If all values in the set are not greater than `val` an empty iterator is returned.
 378     pub fn upper_bound<'a>(&'a self, val: uint) -> SetItems<'a> {
 379         SetItems{iter: self.map.upper_bound(val)}
 380     }
 381 }
 382
 383 impl FromIterator<uint> for TrieSet {
 384     fn from_iter<Iter: Iterator<uint>>(iter: Iter) -> TrieSet {
 385         let mut set = TrieSet::new();
 386         set.extend(iter);
 387         set
 388     }
 389 }
 390
 391 impl Extendable<uint> for TrieSet {
 392     fn extend<Iter: Iterator<uint>>(&mut self, mut iter: Iter) {
 393         for elem in iter {
 394             self.insert(elem);
 395         }
 396     }
 397 }
 398
 399 struct TrieNode<T> {
 400     count: uint,
 401     children: [Child<T>, ..SIZE]
 402 }
 403
 404 impl<T> TrieNode<T> {
 405     #[inline]
 406     fn new() -> TrieNode<T> {
 407         // FIXME: #5244: [Nothing, ..SIZE] should be possible without implicit
 408         // copyability
 409         TrieNode{count: 0,
 410                  children: [Nothing, Nothing, Nothing, Nothing,
 411                             Nothing, Nothing, Nothing, Nothing,
 412                             Nothing, Nothing, Nothing, Nothing,
 413                             Nothing, Nothing, Nothing, Nothing]}
 414     }
 415 }
 416
 417 impl<T> TrieNode<T> {
 418     fn each_reverse<'a>(&'a self, f: |&uint, &'a T| -> bool) -> bool {
 419         for elt in self.children.iter().rev() {
 420             match *elt {
 421                 Internal(ref x) => if !x.each_reverse(|i,t| f(i,t)) { return false },
 422                 External(k, ref v) => if !f(&k, v) { return false },
 423                 Nothing => ()
 424             }
 425         }
 426         true
 427     }
 428 }
 429
 430 // if this was done via a trait, the key could be generic
 431 #[inline]
 432 fn chunk(n: uint, idx: uint) -> uint {
 433     let sh = uint::BITS - (SHIFT * (idx + 1));
 434     (n >> sh) & MASK
 435 }
 436
 437 fn find_mut<'r, T>(child: &'r mut Child<T>, key: uint, idx: uint) -> Option<&'r mut T> {
 438     match *child {
 439         External(stored, ref mut value) if stored == key => Some(value),
 440         External(..) => None,
 441         Internal(ref mut x) => find_mut(&mut x.children[chunk(key, idx)], key, idx + 1),
 442         Nothing => None
 443     }
 444 }
 445
 446 fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
 447              idx: uint) -> Option<T> {
 448     // we branch twice to avoid having to do the `replace` when we
 449     // don't need to; this is much faster, especially for keys that
 450     // have long shared prefixes.
 451     match *child {
 452         Nothing => {
 453             *count += 1;
 454             *child = External(key, value);
 455             return None;
 456         }
 457         Internal(ref mut x) => {
 458             return insert(&mut x.count, &mut x.children[chunk(key, idx)], key, value, idx + 1);
 459         }
 460         External(stored_key, ref mut stored_value) if stored_key == key => {
 461             // swap in the new value and return the old.
 462             return Some(mem::replace(stored_value, value));
 463         }
 464         _ => {}
 465     }
 466
 467     // conflict, an external node with differing keys: we have to
 468     // split the node, so we need the old value by value; hence we
 469     // have to move out of `child`.
 470     match mem::replace(child, Nothing) {
 471         External(stored_key, stored_value) => {
 472             let mut new = box TrieNode::new();
 473             insert(&mut new.count,
 474                    &mut new.children[chunk(stored_key, idx)],
 475                    stored_key, stored_value, idx + 1);
 476             let ret = insert(&mut new.count, &mut new.children[chunk(key, idx)],
 477                          key, value, idx + 1);
 478             *child = Internal(new);
 479             return ret;
 480         }
 481         _ => fail!("unreachable code"),
 482     }
 483 }
 484
 485 fn remove<T>(count: &mut uint, child: &mut Child<T>, key: uint,
 486              idx: uint) -> Option<T> {
 487     let (ret, this) = match *child {
 488       External(stored, _) if stored == key => {
 489         match mem::replace(child, Nothing) {
 490             External(_, value) => (Some(value), true),
 491             _ => fail!()
 492         }
 493       }
 494       External(..) => (None, false),
 495       Internal(ref mut x) => {
 496           let ret = remove(&mut x.count, &mut x.children[chunk(key, idx)],
 497                            key, idx + 1);
 498           (ret, x.count == 0)
 499       }
 500       Nothing => (None, false)
 501     };
 502
 503     if this {
 504         *child = Nothing;
 505         *count -= 1;
 506     }
 507     return ret;
 508 }
 509
 510 /// Forward iterator over a map
 511 pub struct Entries<'a, T> {
 512     stack: [slice::Items<'a, Child<T>>, .. NUM_CHUNKS],
 513     length: uint,
 514     remaining_min: uint,
 515     remaining_max: uint
 516 }
 517
 518 /// Forward iterator over the key-value pairs of a map, with the
 519 /// values being mutable.
 520 pub struct MutEntries<'a, T> {
 521     stack: [slice::MutItems<'a, Child<T>>, .. NUM_CHUNKS],
 522     length: uint,
 523     remaining_min: uint,
 524     remaining_max: uint
 525 }
 526
 527 // FIXME #5846: see `addr!` above.
 528 macro_rules! item { ($i:item) => {$i}}
 529
 530 macro_rules! iterator_impl {
 531     ($name:ident,
 532      iter = $iter:ident,
 533      mutability = $($mut_:tt)*) => {
 534         impl<'a, T> $name<'a, T> {
 535             // Create new zero'd iterator. We have a thin gilding of safety by
 536             // using init rather than uninit, so that the worst that can happen
 537             // from failing to initialise correctly after calling these is a
 538             // segfault.
 539             #[cfg(target_word_size="32")]
 540             unsafe fn new() -> $name<'a, T> {
 541                 $name {
 542                     remaining_min: 0,
 543                     remaining_max: 0,
 544                     length: 0,
 545                     // ick :( ... at least the compiler will tell us if we screwed up.
 546                     stack: [zeroed(), zeroed(), zeroed(), zeroed(), zeroed(),
 547                             zeroed(), zeroed(), zeroed()]
 548                 }
 549             }
 550
 551             #[cfg(target_word_size="64")]
 552             unsafe fn new() -> $name<'a, T> {
 553                 $name {
 554                     remaining_min: 0,
 555                     remaining_max: 0,
 556                     length: 0,
 557                     stack: [zeroed(), zeroed(), zeroed(), zeroed(),
 558                             zeroed(), zeroed(), zeroed(), zeroed(),
 559                             zeroed(), zeroed(), zeroed(), zeroed(),
 560                             zeroed(), zeroed(), zeroed(), zeroed()]
 561                 }
 562             }
 563         }
 564
 565         item!(impl<'a, T> Iterator<(uint, &'a $($mut_)* T)> for $name<'a, T> {
 566                 // you might wonder why we're not even trying to act within the
 567                 // rules, and are just manipulating raw pointers like there's no
 568                 // such thing as invalid pointers and memory unsafety. The
 569                 // reason is performance, without doing this we can get the
 570                 // bench_iter_large microbenchmark down to about 30000 ns/iter
 571                 // (using .unsafe_ref to index self.stack directly, 38000
 572                 // ns/iter with [] checked indexing), but this smashes that down
 573                 // to 13500 ns/iter.
 574                 //
 575                 // Fortunately, it's still safe...
 576                 //
 577                 // We have an invariant that every Internal node
 578                 // corresponds to one push to self.stack, and one pop,
 579                 // nested appropriately. self.stack has enough storage
 580                 // to store the maximum depth of Internal nodes in the
 581                 // trie (8 on 32-bit platforms, 16 on 64-bit).
 582                 fn next(&mut self) -> Option<(uint, &'a $($mut_)* T)> {
 583                     let start_ptr = self.stack.as_mut_ptr();
 584
 585                     unsafe {
 586                         // write_ptr is the next place to write to the stack.
 587                         // invariant: start_ptr <= write_ptr < end of the
 588                         // vector.
 589                         let mut write_ptr = start_ptr.offset(self.length as int);
 590                         while write_ptr != start_ptr {
 591                             // indexing back one is safe, since write_ptr >
 592                             // start_ptr now.
 593                             match (*write_ptr.offset(-1)).next() {
 594                                 // exhausted this iterator (i.e. finished this
 595                                 // Internal node), so pop from the stack.
 596                                 //
 597                                 // don't bother clearing the memory, because the
 598                                 // next time we use it we'll've written to it
 599                                 // first.
 600                                 None => write_ptr = write_ptr.offset(-1),
 601                                 Some(child) => {
 602                                     addr!(match *child {
 603                                             Internal(ref $($mut_)* node) => {
 604                                                 // going down a level, so push
 605                                                 // to the stack (this is the
 606                                                 // write referenced above)
 607                                                 *write_ptr = node.children.$iter();
 608                                                 write_ptr = write_ptr.offset(1);
 609                                             }
 610                                             External(key, ref $($mut_)* value) => {
 611                                                 self.remaining_max -= 1;
 612                                                 if self.remaining_min > 0 {
 613                                                     self.remaining_min -= 1;
 614                                                 }
 615                                                 // store the new length of the
 616                                                 // stack, based on our current
 617                                                 // position.
 618                                                 self.length = (write_ptr as uint
 619                                                                - start_ptr as uint) /
 620                                                     mem::size_of_val(&*write_ptr);
 621
 622                                                 return Some((key, value));
 623                                             }
 624                                             Nothing => {}
 625                                         })
 626                                 }
 627                             }
 628                         }
 629                     }
 630                     return None;
 631                 }
 632
 633                 #[inline]
 634                 fn size_hint(&self) -> (uint, Option<uint>) {
 635                     (self.remaining_min, Some(self.remaining_max))
 636                 }
 637             })
 638     }
 639 }
 640
 641 iterator_impl! { Entries, iter = iter, mutability = }
 642 iterator_impl! { MutEntries, iter = mut_iter, mutability = mut }
 643
 644 /// Forward iterator over a set
 645 pub struct SetItems<'a> {
 646     iter: Entries<'a, ()>
 647 }
 648
 649 impl<'a> Iterator<uint> for SetItems<'a> {
 650     fn next(&mut self) -> Option<uint> {
 651         self.iter.next().map(|(key, _)| key)
 652     }
 653
 654     fn size_hint(&self) -> (uint, Option<uint>) {
 655         self.iter.size_hint()
 656     }
 657 }
 658
 659 #[cfg(test)]
 660 mod test_map {
 661     use std::prelude::*;
 662     use std::iter::range_step;
 663     use std::uint;
 664
 665     use {MutableMap, Map};
 666     use super::{TrieMap, TrieNode, Internal, External, Nothing};
 667
 668     fn check_integrity<T>(trie: &TrieNode<T>) {
 669         assert!(trie.count != 0);
 670
 671         let mut sum = 0;
 672
 673         for x in trie.children.iter() {
 674             match *x {
 675               Nothing => (),
 676               Internal(ref y) => {
 677                   check_integrity(&**y);
 678                   sum += 1
 679               }
 680               External(_, _) => { sum += 1 }
 681             }
 682         }
 683
 684         assert_eq!(sum, trie.count);
 685     }
 686
 687     #[test]
 688     fn test_find_mut() {
 689         let mut m = TrieMap::new();
 690         assert!(m.insert(1u, 12i));
 691         assert!(m.insert(2u, 8i));
 692         assert!(m.insert(5u, 14i));
 693         let new = 100;
 694         match m.find_mut(&5) {
 695             None => fail!(), Some(x) => *x = new
 696         }
 697         assert_eq!(m.find(&5), Some(&new));
 698     }
 699
 700     #[test]
 701     fn test_find_mut_missing() {
 702         let mut m = TrieMap::new();
 703         assert!(m.find_mut(&0).is_none());
 704         assert!(m.insert(1u, 12i));
 705         assert!(m.find_mut(&0).is_none());
 706         assert!(m.insert(2, 8));
 707         assert!(m.find_mut(&0).is_none());
 708     }
 709
 710     #[test]
 711     fn test_step() {
 712         let mut trie = TrieMap::new();
 713         let n = 300u;
 714
 715         for x in range_step(1u, n, 2) {
 716             assert!(trie.insert(x, x + 1));
 717             assert!(trie.contains_key(&x));
 718             check_integrity(&trie.root);
 719         }
 720
 721         for x in range_step(0u, n, 2) {
 722             assert!(!trie.contains_key(&x));
 723             assert!(trie.insert(x, x + 1));
 724             check_integrity(&trie.root);
 725         }
 726
 727         for x in range(0u, n) {
 728             assert!(trie.contains_key(&x));
 729             assert!(!trie.insert(x, x + 1));
 730             check_integrity(&trie.root);
 731         }
 732
 733         for x in range_step(1u, n, 2) {
 734             assert!(trie.remove(&x));
 735             assert!(!trie.contains_key(&x));
 736             check_integrity(&trie.root);
 737         }
 738
 739         for x in range_step(0u, n, 2) {
 740             assert!(trie.contains_key(&x));
 741             assert!(!trie.insert(x, x + 1));
 742             check_integrity(&trie.root);
 743         }
 744     }
 745
 746     #[test]
 747     fn test_each_reverse() {
 748         let mut m = TrieMap::new();
 749
 750         assert!(m.insert(3, 6));
 751         assert!(m.insert(0, 0));
 752         assert!(m.insert(4, 8));
 753         assert!(m.insert(2, 4));
 754         assert!(m.insert(1, 2));
 755
 756         let mut n = 4;
 757         m.each_reverse(|k, v| {
 758             assert_eq!(*k, n);
 759             assert_eq!(*v, n * 2);
 760             n -= 1;
 761             true
 762         });
 763     }
 764
 765     #[test]
 766     fn test_each_reverse_break() {
 767         let mut m = TrieMap::new();
 768
 769         for x in range(uint::MAX - 10000, uint::MAX).rev() {
 770             m.insert(x, x / 2);
 771         }
 772
 773         let mut n = uint::MAX - 1;
 774         m.each_reverse(|k, v| {
 775             if n == uint::MAX - 5000 { false } else {
 776                 assert!(n > uint::MAX - 5000);
 777
 778                 assert_eq!(*k, n);
 779                 assert_eq!(*v, n / 2);
 780                 n -= 1;
 781                 true
 782             }
 783         });
 784     }
 785
 786     #[test]
 787     fn test_swap() {
 788         let mut m = TrieMap::new();
 789         assert_eq!(m.swap(1u, 2i), None);
 790         assert_eq!(m.swap(1u, 3i), Some(2));
 791         assert_eq!(m.swap(1u, 4i), Some(3));
 792     }
 793
 794     #[test]
 795     fn test_pop() {
 796         let mut m = TrieMap::new();
 797         m.insert(1u, 2i);
 798         assert_eq!(m.pop(&1), Some(2));
 799         assert_eq!(m.pop(&1), None);
 800     }
 801
 802     #[test]
 803     fn test_from_iter() {
 804         let xs = vec![(1u, 1i), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
 805
 806         let map: TrieMap<int> = xs.iter().map(|&x| x).collect();
 807
 808         for &(k, v) in xs.iter() {
 809             assert_eq!(map.find(&k), Some(&v));
 810         }
 811     }
 812
 813     #[test]
 814     fn test_iteration() {
 815         let empty_map : TrieMap<uint> = TrieMap::new();
 816         assert_eq!(empty_map.iter().next(), None);
 817
 818         let first = uint::MAX - 10000;
 819         let last = uint::MAX;
 820
 821         let mut map = TrieMap::new();
 822         for x in range(first, last).rev() {
 823             map.insert(x, x / 2);
 824         }
 825
 826         let mut i = 0;
 827         for (k, &v) in map.iter() {
 828             assert_eq!(k, first + i);
 829             assert_eq!(v, k / 2);
 830             i += 1;
 831         }
 832         assert_eq!(i, last - first);
 833     }
 834
 835     #[test]
 836     fn test_mut_iter() {
 837         let mut empty_map : TrieMap<uint> = TrieMap::new();
 838         assert!(empty_map.mut_iter().next().is_none());
 839
 840         let first = uint::MAX - 10000;
 841         let last = uint::MAX;
 842
 843         let mut map = TrieMap::new();
 844         for x in range(first, last).rev() {
 845             map.insert(x, x / 2);
 846         }
 847
 848         let mut i = 0;
 849         for (k, v) in map.mut_iter() {
 850             assert_eq!(k, first + i);
 851             *v -= k / 2;
 852             i += 1;
 853         }
 854         assert_eq!(i, last - first);
 855
 856         assert!(map.iter().all(|(_, &v)| v == 0));
 857     }
 858
 859     #[test]
 860     fn test_bound() {
 861         let empty_map : TrieMap<uint> = TrieMap::new();
 862         assert_eq!(empty_map.lower_bound(0).next(), None);
 863         assert_eq!(empty_map.upper_bound(0).next(), None);
 864
 865         let last = 999u;
 866         let step = 3u;
 867         let value = 42u;
 868
 869         let mut map : TrieMap<uint> = TrieMap::new();
 870         for x in range_step(0u, last, step) {
 871             assert!(x % step == 0);
 872             map.insert(x, value);
 873         }
 874
 875         for i in range(0u, last - step) {
 876             let mut lb = map.lower_bound(i);
 877             let mut ub = map.upper_bound(i);
 878             let next_key = i - i % step + step;
 879             let next_pair = (next_key, &value);
 880             if i % step == 0 {
 881                 assert_eq!(lb.next(), Some((i, &value)));
 882             } else {
 883                 assert_eq!(lb.next(), Some(next_pair));
 884             }
 885             assert_eq!(ub.next(), Some(next_pair));
 886         }
 887
 888         let mut lb = map.lower_bound(last - step);
 889         assert_eq!(lb.next(), Some((last - step, &value)));
 890         let mut ub = map.upper_bound(last - step);
 891         assert_eq!(ub.next(), None);
 892
 893         for i in range(last - step + 1, last) {
 894             let mut lb = map.lower_bound(i);
 895             assert_eq!(lb.next(), None);
 896             let mut ub = map.upper_bound(i);
 897             assert_eq!(ub.next(), None);
 898         }
 899     }
 900
 901     #[test]
 902     fn test_mut_bound() {
 903         let empty_map : TrieMap<uint> = TrieMap::new();
 904         assert_eq!(empty_map.lower_bound(0).next(), None);
 905         assert_eq!(empty_map.upper_bound(0).next(), None);
 906
 907         let mut m_lower = TrieMap::new();
 908         let mut m_upper = TrieMap::new();
 909         for i in range(0u, 100) {
 910             m_lower.insert(2 * i, 4 * i);
 911             m_upper.insert(2 * i, 4 * i);
 912         }
 913
 914         for i in range(0u, 199) {
 915             let mut lb_it = m_lower.mut_lower_bound(i);
 916             let (k, v) = lb_it.next().unwrap();
 917             let lb = i + i % 2;
 918             assert_eq!(lb, k);
 919             *v -= k;
 920         }
 921
 922         for i in range(0u, 198) {
 923             let mut ub_it = m_upper.mut_upper_bound(i);
 924             let (k, v) = ub_it.next().unwrap();
 925             let ub = i + 2 - i % 2;
 926             assert_eq!(ub, k);
 927             *v -= k;
 928         }
 929
 930         assert!(m_lower.mut_lower_bound(199).next().is_none());
 931         assert!(m_upper.mut_upper_bound(198).next().is_none());
 932
 933         assert!(m_lower.iter().all(|(_, &x)| x == 0));
 934         assert!(m_upper.iter().all(|(_, &x)| x == 0));
 935     }
 936 }
 937
 938 #[cfg(test)]
 939 mod bench_map {
 940     use std::prelude::*;
 941     use std::rand::{weak_rng, Rng};
 942     use test::Bencher;
 943
 944     use MutableMap;
 945     use super::TrieMap;
 946
 947     #[bench]
 948     fn bench_iter_small(b: &mut Bencher) {
 949         let mut m = TrieMap::<uint>::new();
 950         let mut rng = weak_rng();
 951         for _ in range(0u, 20) {
 952             m.insert(rng.gen(), rng.gen());
 953         }
 954
 955         b.iter(|| for _ in m.iter() {})
 956     }
 957
 958     #[bench]
 959     fn bench_iter_large(b: &mut Bencher) {
 960         let mut m = TrieMap::<uint>::new();
 961         let mut rng = weak_rng();
 962         for _ in range(0u, 1000) {
 963             m.insert(rng.gen(), rng.gen());
 964         }
 965
 966         b.iter(|| for _ in m.iter() {})
 967     }
 968
 969     #[bench]
 970     fn bench_lower_bound(b: &mut Bencher) {
 971         let mut m = TrieMap::<uint>::new();
 972         let mut rng = weak_rng();
 973         for _ in range(0u, 1000) {
 974             m.insert(rng.gen(), rng.gen());
 975         }
 976
 977         b.iter(|| {
 978                 for _ in range(0u, 10) {
 979                     m.lower_bound(rng.gen());
 980                 }
 981             });
 982     }
 983
 984     #[bench]
 985     fn bench_upper_bound(b: &mut Bencher) {
 986         let mut m = TrieMap::<uint>::new();
 987         let mut rng = weak_rng();
 988         for _ in range(0u, 1000) {
 989             m.insert(rng.gen(), rng.gen());
 990         }
 991
 992         b.iter(|| {
 993                 for _ in range(0u, 10) {
 994                     m.upper_bound(rng.gen());
 995                 }
 996             });
 997     }
 998
 999     #[bench]
1000     fn bench_insert_large(b: &mut Bencher) {
1001         let mut m = TrieMap::<[uint, .. 10]>::new();
1002         let mut rng = weak_rng();
1003
1004         b.iter(|| {
1005                 for _ in range(0u, 1000) {
1006                     m.insert(rng.gen(), [1, .. 10]);
1007                 }
1008             })
1009     }
1010     #[bench]
1011     fn bench_insert_large_low_bits(b: &mut Bencher) {
1012         let mut m = TrieMap::<[uint, .. 10]>::new();
1013         let mut rng = weak_rng();
1014
1015         b.iter(|| {
1016                 for _ in range(0u, 1000) {
1017                     // only have the last few bits set.
1018                     m.insert(rng.gen::<uint>() & 0xff_ff, [1, .. 10]);
1019                 }
1020             })
1021     }
1022
1023     #[bench]
1024     fn bench_insert_small(b: &mut Bencher) {
1025         let mut m = TrieMap::<()>::new();
1026         let mut rng = weak_rng();
1027
1028         b.iter(|| {
1029                 for _ in range(0u, 1000) {
1030                     m.insert(rng.gen(), ());
1031                 }
1032             })
1033     }
1034     #[bench]
1035     fn bench_insert_small_low_bits(b: &mut Bencher) {
1036         let mut m = TrieMap::<()>::new();
1037         let mut rng = weak_rng();
1038
1039         b.iter(|| {
1040                 for _ in range(0u, 1000) {
1041                     // only have the last few bits set.
1042                     m.insert(rng.gen::<uint>() & 0xff_ff, ());
1043                 }
1044             })
1045     }
1046 }
1047
1048 #[cfg(test)]
1049 mod test_set {
1050     use std::prelude::*;
1051     use std::uint;
1052
1053     use {MutableSet, Set};
1054     use super::TrieSet;
1055
1056     #[test]
1057     fn test_sane_chunk() {
1058         let x = 1;
1059         let y = 1 << (uint::BITS - 1);
1060
1061         let mut trie = TrieSet::new();
1062
1063         assert!(trie.insert(x));
1064         assert!(trie.insert(y));
1065
1066         assert_eq!(trie.len(), 2);
1067
1068         let expected = [x, y];
1069
1070         for (i, x) in trie.iter().enumerate() {
1071             assert_eq!(expected[i], x);
1072         }
1073     }
1074
1075     #[test]
1076     fn test_from_iter() {
1077         let xs = vec![9u, 8, 7, 6, 5, 4, 3, 2, 1];
1078
1079         let set: TrieSet = xs.iter().map(|&x| x).collect();
1080
1081         for x in xs.iter() {
1082             assert!(set.contains(x));
1083         }
1084     }
1085 }