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