src/libstd/trie.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 for integer keys implemented as a radix trie
  12
  13 use prelude::*;
  14 use iterator::{IteratorUtil, FromIterator, Extendable};
  15 use uint;
  16 use util::{swap, replace};
  17 use vec;
  18
  19 // FIXME: #5244: need to manually update the TrieNode constructor
  20 static SHIFT: uint = 4;
  21 static SIZE: uint = 1 << SHIFT;
  22 static MASK: uint = SIZE - 1;
  23
  24 enum Child<T> {
  25     Internal(~TrieNode<T>),
  26     External(uint, T),
  27     Nothing
  28 }
  29
  30 #[allow(missing_doc)]
  31 pub struct TrieMap<T> {
  32     priv root: TrieNode<T>,
  33     priv length: uint
  34 }
  35
  36 impl<T> Container for TrieMap<T> {
  37     /// Return the number of elements in the map
  38     #[inline]
  39     fn len(&self) -> uint { self.length }
  40 }
  41
  42 impl<T> Mutable for TrieMap<T> {
  43     /// Clear the map, removing all values.
  44     #[inline]
  45     fn clear(&mut self) {
  46         self.root = TrieNode::new();
  47         self.length = 0;
  48     }
  49 }
  50
  51 impl<T> Map<uint, T> for TrieMap<T> {
  52     /// Return a reference to the value corresponding to the key
  53     #[inline]
  54     fn find<'a>(&'a self, key: &uint) -> Option<&'a T> {
  55         let mut node: &'a TrieNode<T> = &self.root;
  56         let mut idx = 0;
  57         loop {
  58             match node.children[chunk(*key, idx)] {
  59               Internal(ref x) => node = &**x,
  60               External(stored, ref value) => {
  61                 if stored == *key {
  62                     return Some(value)
  63                 } else {
  64                     return None
  65                 }
  66               }
  67               Nothing => return None
  68             }
  69             idx += 1;
  70         }
  71     }
  72 }
  73
  74 impl<T> MutableMap<uint, T> for TrieMap<T> {
  75     /// Return a mutable reference to the value corresponding to the key
  76     #[inline]
  77     fn find_mut<'a>(&'a mut self, key: &uint) -> Option<&'a mut T> {
  78         find_mut(&mut self.root.children[chunk(*key, 0)], *key, 1)
  79     }
  80
  81     /// Insert a key-value pair from the map. If the key already had a value
  82     /// present in the map, that value is returned. Otherwise None is returned.
  83     fn swap(&mut self, key: uint, value: T) -> Option<T> {
  84         let ret = insert(&mut self.root.count,
  85                          &mut self.root.children[chunk(key, 0)],
  86                          key, value, 1);
  87         if ret.is_none() { self.length += 1 }
  88         ret
  89     }
  90
  91     /// Removes a key from the map, returning the value at the key if the key
  92     /// was previously in the map.
  93     fn pop(&mut self, key: &uint) -> Option<T> {
  94         let ret = remove(&mut self.root.count,
  95                          &mut self.root.children[chunk(*key, 0)],
  96                          *key, 1);
  97         if ret.is_some() { self.length -= 1 }
  98         ret
  99     }
 100 }
 101
 102 impl<T> TrieMap<T> {
 103     /// Create an empty TrieMap
 104     #[inline]
 105     pub fn new() -> TrieMap<T> {
 106         TrieMap{root: TrieNode::new(), length: 0}
 107     }
 108
 109     /// Visit all key-value pairs in reverse order
 110     #[inline]
 111     pub fn each_reverse<'a>(&'a self, f: &fn(&uint, &'a T) -> bool) -> bool {
 112         self.root.each_reverse(f)
 113     }
 114
 115     /// Visit all key-value pairs in order
 116     #[inline]
 117     pub fn each<'a>(&'a self, f: &fn(&uint, &'a T) -> bool) -> bool {
 118         self.root.each(f)
 119     }
 120
 121     /// Visit all keys in order
 122     #[inline]
 123     pub fn each_key(&self, f: &fn(&uint) -> bool) -> bool {
 124         self.each(|k, _| f(k))
 125     }
 126
 127     /// Visit all values in order
 128     #[inline]
 129     pub fn each_value<'a>(&'a self, f: &fn(&'a T) -> bool) -> bool {
 130         self.each(|_, v| f(v))
 131     }
 132
 133     /// Iterate over the map and mutate the contained values
 134     #[inline]
 135     pub fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) -> bool {
 136         self.root.mutate_values(f)
 137     }
 138
 139     /// Visit all keys in reverse order
 140     #[inline]
 141     pub fn each_key_reverse(&self, f: &fn(&uint) -> bool) -> bool {
 142         self.each_reverse(|k, _| f(k))
 143     }
 144
 145     /// Visit all values in reverse order
 146     #[inline]
 147     pub fn each_value_reverse(&self, f: &fn(&T) -> bool) -> bool {
 148         self.each_reverse(|_, v| f(v))
 149     }
 150
 151     /// Get an iterator over the key-value pairs in the map
 152     pub fn iter<'a>(&'a self) -> TrieMapIterator<'a, T> {
 153         TrieMapIterator {
 154             stack: ~[self.root.children.iter()],
 155             remaining_min: self.length,
 156             remaining_max: self.length
 157         }
 158     }
 159 }
 160
 161 impl<T, Iter: Iterator<(uint, T)>> FromIterator<(uint, T), Iter> for TrieMap<T> {
 162     fn from_iterator(iter: &mut Iter) -> TrieMap<T> {
 163         let mut map = TrieMap::new();
 164         map.extend(iter);
 165         map
 166     }
 167 }
 168
 169 impl<T, Iter: Iterator<(uint, T)>> Extendable<(uint, T), Iter> for TrieMap<T> {
 170     fn extend(&mut self, iter: &mut Iter) {
 171         for (k, v) in *iter {
 172             self.insert(k, v);
 173         }
 174     }
 175 }
 176
 177 #[allow(missing_doc)]
 178 pub struct TrieSet {
 179     priv map: TrieMap<()>
 180 }
 181
 182 impl Container for TrieSet {
 183     /// Return the number of elements in the set
 184     #[inline]
 185     fn len(&self) -> uint { self.map.len() }
 186 }
 187
 188 impl Mutable for TrieSet {
 189     /// Clear the set, removing all values.
 190     #[inline]
 191     fn clear(&mut self) { self.map.clear() }
 192 }
 193
 194 impl TrieSet {
 195     /// Create an empty TrieSet
 196     #[inline]
 197     pub fn new() -> TrieSet {
 198         TrieSet{map: TrieMap::new()}
 199     }
 200
 201     /// Return true if the set contains a value
 202     #[inline]
 203     pub fn contains(&self, value: &uint) -> bool {
 204         self.map.contains_key(value)
 205     }
 206
 207     /// Add a value to the set. Return true if the value was not already
 208     /// present in the set.
 209     #[inline]
 210     pub fn insert(&mut self, value: uint) -> bool {
 211         self.map.insert(value, ())
 212     }
 213
 214     /// Remove a value from the set. Return true if the value was
 215     /// present in the set.
 216     #[inline]
 217     pub fn remove(&mut self, value: &uint) -> bool {
 218         self.map.remove(value)
 219     }
 220
 221     /// Visit all values in order
 222     #[inline]
 223     pub fn each(&self, f: &fn(&uint) -> bool) -> bool { self.map.each_key(f) }
 224
 225     /// Visit all values in reverse order
 226     #[inline]
 227     pub fn each_reverse(&self, f: &fn(&uint) -> bool) -> bool {
 228         self.map.each_key_reverse(f)
 229     }
 230
 231     /// Get an iterator over the values in the set
 232     #[inline]
 233     pub fn iter<'a>(&'a self) -> TrieSetIterator<'a> {
 234         TrieSetIterator{iter: self.map.iter()}
 235     }
 236 }
 237
 238 impl<Iter: Iterator<uint>> FromIterator<uint, Iter> for TrieSet {
 239     fn from_iterator(iter: &mut Iter) -> TrieSet {
 240         let mut set = TrieSet::new();
 241         set.extend(iter);
 242         set
 243     }
 244 }
 245
 246 impl<Iter: Iterator<uint>> Extendable<uint, Iter> for TrieSet {
 247     fn extend(&mut self, iter: &mut Iter) {
 248         for elem in *iter {
 249             self.insert(elem);
 250         }
 251     }
 252 }
 253
 254 struct TrieNode<T> {
 255     count: uint,
 256     children: [Child<T>, ..SIZE]
 257 }
 258
 259 impl<T> TrieNode<T> {
 260     #[inline]
 261     fn new() -> TrieNode<T> {
 262         // FIXME: #5244: [Nothing, ..SIZE] should be possible without implicit
 263         // copyability
 264         TrieNode{count: 0,
 265                  children: [Nothing, Nothing, Nothing, Nothing,
 266                             Nothing, Nothing, Nothing, Nothing,
 267                             Nothing, Nothing, Nothing, Nothing,
 268                             Nothing, Nothing, Nothing, Nothing]}
 269     }
 270 }
 271
 272 impl<T> TrieNode<T> {
 273     fn each<'a>(&'a self, f: &fn(&uint, &'a T) -> bool) -> bool {
 274         for idx in range(0u, self.children.len()) {
 275             match self.children[idx] {
 276                 Internal(ref x) => if !x.each(|i,t| f(i,t)) { return false },
 277                 External(k, ref v) => if !f(&k, v) { return false },
 278                 Nothing => ()
 279             }
 280         }
 281         true
 282     }
 283
 284     fn each_reverse<'a>(&'a self, f: &fn(&uint, &'a T) -> bool) -> bool {
 285         do uint::range_rev(self.children.len(), 0) |idx| {
 286             match self.children[idx] {
 287                 Internal(ref x) => x.each_reverse(|i,t| f(i,t)),
 288                 External(k, ref v) => f(&k, v),
 289                 Nothing => true
 290             }
 291         }
 292     }
 293
 294     fn mutate_values<'a>(&'a mut self, f: &fn(&uint, &mut T) -> bool) -> bool {
 295         for child in self.children.mut_iter() {
 296             match *child {
 297                 Internal(ref mut x) => if !x.mutate_values(|i,t| f(i,t)) {
 298                     return false
 299                 },
 300                 External(k, ref mut v) => if !f(&k, v) { return false },
 301                 Nothing => ()
 302             }
 303         }
 304         true
 305     }
 306 }
 307
 308 // if this was done via a trait, the key could be generic
 309 #[inline]
 310 fn chunk(n: uint, idx: uint) -> uint {
 311     let sh = uint::bits - (SHIFT * (idx + 1));
 312     (n >> sh) & MASK
 313 }
 314
 315 fn find_mut<'r, T>(child: &'r mut Child<T>, key: uint, idx: uint) -> Option<&'r mut T> {
 316     match *child {
 317         External(_, ref mut value) => Some(value),
 318         Internal(ref mut x) => find_mut(&mut x.children[chunk(key, idx)], key, idx + 1),
 319         Nothing => None
 320     }
 321 }
 322
 323 fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
 324              idx: uint) -> Option<T> {
 325     let mut tmp = Nothing;
 326     let ret;
 327     swap(&mut tmp, child);
 328
 329     *child = match tmp {
 330       External(stored_key, stored_value) => {
 331           if stored_key == key {
 332               ret = Some(stored_value);
 333               External(stored_key, value)
 334           } else {
 335               // conflict - split the node
 336               let mut new = ~TrieNode::new();
 337               insert(&mut new.count,
 338                      &mut new.children[chunk(stored_key, idx)],
 339                      stored_key, stored_value, idx + 1);
 340               ret = insert(&mut new.count, &mut new.children[chunk(key, idx)],
 341                            key, value, idx + 1);
 342               Internal(new)
 343           }
 344       }
 345       Internal(x) => {
 346         let mut x = x;
 347         ret = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
 348                      value, idx + 1);
 349         Internal(x)
 350       }
 351       Nothing => {
 352         *count += 1;
 353         ret = None;
 354         External(key, value)
 355       }
 356     };
 357     return ret;
 358 }
 359
 360 fn remove<T>(count: &mut uint, child: &mut Child<T>, key: uint,
 361              idx: uint) -> Option<T> {
 362     let (ret, this) = match *child {
 363       External(stored, _) if stored == key => {
 364         match replace(child, Nothing) {
 365             External(_, value) => (Some(value), true),
 366             _ => fail!()
 367         }
 368       }
 369       External(*) => (None, false),
 370       Internal(ref mut x) => {
 371           let ret = remove(&mut x.count, &mut x.children[chunk(key, idx)],
 372                            key, idx + 1);
 373           (ret, x.count == 0)
 374       }
 375       Nothing => (None, false)
 376     };
 377
 378     if this {
 379         *child = Nothing;
 380         *count -= 1;
 381     }
 382     return ret;
 383 }
 384
 385 /// Forward iterator over a map
 386 pub struct TrieMapIterator<'self, T> {
 387     priv stack: ~[vec::VecIterator<'self, Child<T>>],
 388     priv remaining_min: uint,
 389     priv remaining_max: uint
 390 }
 391
 392 impl<'self, T> Iterator<(uint, &'self T)> for TrieMapIterator<'self, T> {
 393     fn next(&mut self) -> Option<(uint, &'self T)> {
 394         while !self.stack.is_empty() {
 395             match self.stack[self.stack.len() - 1].next() {
 396                 None => {
 397                     self.stack.pop();
 398                 }
 399                 Some(ref child) => {
 400                     match **child {
 401                         Internal(ref node) => {
 402                             self.stack.push(node.children.iter());
 403                         }
 404                         External(key, ref value) => {
 405                             self.remaining_max -= 1;
 406                             if self.remaining_min > 0 {
 407                                 self.remaining_min -= 1;
 408                             }
 409                             return Some((key, value));
 410                         }
 411                         Nothing => {}
 412                     }
 413                 }
 414             }
 415         }
 416         return None;
 417     }
 418
 419     #[inline]
 420     fn size_hint(&self) -> (uint, Option<uint>) {
 421         (self.remaining_min, Some(self.remaining_max))
 422     }
 423 }
 424
 425 /// Forward iterator over a set
 426 pub struct TrieSetIterator<'self> {
 427     priv iter: TrieMapIterator<'self, ()>
 428 }
 429
 430 impl<'self> Iterator<uint> for TrieSetIterator<'self> {
 431     fn next(&mut self) -> Option<uint> {
 432         do self.iter.next().map |&(key, _)| { key }
 433     }
 434
 435     fn size_hint(&self) -> (uint, Option<uint>) {
 436         self.iter.size_hint()
 437     }
 438 }
 439
 440 #[cfg(test)]
 441 pub fn check_integrity<T>(trie: &TrieNode<T>) {
 442     assert!(trie.count != 0);
 443
 444     let mut sum = 0;
 445
 446     for x in trie.children.iter() {
 447         match *x {
 448           Nothing => (),
 449           Internal(ref y) => {
 450               check_integrity(&**y);
 451               sum += 1
 452           }
 453           External(_, _) => { sum += 1 }
 454         }
 455     }
 456
 457     assert_eq!(sum, trie.count);
 458 }
 459
 460 #[cfg(test)]
 461 mod test_map {
 462     use super::*;
 463     use prelude::*;
 464     use uint;
 465
 466     #[test]
 467     fn test_find_mut() {
 468         let mut m = TrieMap::new();
 469         assert!(m.insert(1, 12));
 470         assert!(m.insert(2, 8));
 471         assert!(m.insert(5, 14));
 472         let new = 100;
 473         match m.find_mut(&5) {
 474             None => fail!(), Some(x) => *x = new
 475         }
 476         assert_eq!(m.find(&5), Some(&new));
 477     }
 478
 479     #[test]
 480     fn test_step() {
 481         let mut trie = TrieMap::new();
 482         let n = 300;
 483
 484         do uint::range_step(1, n, 2) |x| {
 485             assert!(trie.insert(x, x + 1));
 486             assert!(trie.contains_key(&x));
 487             check_integrity(&trie.root);
 488             true
 489         };
 490
 491         do uint::range_step(0, n, 2) |x| {
 492             assert!(!trie.contains_key(&x));
 493             assert!(trie.insert(x, x + 1));
 494             check_integrity(&trie.root);
 495             true
 496         };
 497
 498         for x in range(0u, n) {
 499             assert!(trie.contains_key(&x));
 500             assert!(!trie.insert(x, x + 1));
 501             check_integrity(&trie.root);
 502         }
 503
 504         do uint::range_step(1, n, 2) |x| {
 505             assert!(trie.remove(&x));
 506             assert!(!trie.contains_key(&x));
 507             check_integrity(&trie.root);
 508             true
 509         };
 510
 511         do uint::range_step(0, n, 2) |x| {
 512             assert!(trie.contains_key(&x));
 513             assert!(!trie.insert(x, x + 1));
 514             check_integrity(&trie.root);
 515             true
 516         };
 517     }
 518
 519     #[test]
 520     fn test_each() {
 521         let mut m = TrieMap::new();
 522
 523         assert!(m.insert(3, 6));
 524         assert!(m.insert(0, 0));
 525         assert!(m.insert(4, 8));
 526         assert!(m.insert(2, 4));
 527         assert!(m.insert(1, 2));
 528
 529         let mut n = 0;
 530         do m.each |k, v| {
 531             assert_eq!(*k, n);
 532             assert_eq!(*v, n * 2);
 533             n += 1;
 534             true
 535         };
 536     }
 537
 538     #[test]
 539     fn test_each_break() {
 540         let mut m = TrieMap::new();
 541
 542         do uint::range_rev(uint::max_value, uint::max_value - 10000) |x| {
 543             m.insert(x, x / 2);
 544             true
 545         };
 546
 547         let mut n = uint::max_value - 10000;
 548         do m.each |k, v| {
 549             if n == uint::max_value - 5000 { false } else {
 550                 assert!(n < uint::max_value - 5000);
 551
 552                 assert_eq!(*k, n);
 553                 assert_eq!(*v, n / 2);
 554                 n += 1;
 555                 true
 556             }
 557         };
 558     }
 559
 560     #[test]
 561     fn test_each_reverse() {
 562         let mut m = TrieMap::new();
 563
 564         assert!(m.insert(3, 6));
 565         assert!(m.insert(0, 0));
 566         assert!(m.insert(4, 8));
 567         assert!(m.insert(2, 4));
 568         assert!(m.insert(1, 2));
 569
 570         let mut n = 4;
 571         do m.each_reverse |k, v| {
 572             assert_eq!(*k, n);
 573             assert_eq!(*v, n * 2);
 574             n -= 1;
 575             true
 576         };
 577     }
 578
 579     #[test]
 580     fn test_each_reverse_break() {
 581         let mut m = TrieMap::new();
 582
 583         do uint::range_rev(uint::max_value, uint::max_value - 10000) |x| {
 584             m.insert(x, x / 2);
 585             true
 586         };
 587
 588         let mut n = uint::max_value - 1;
 589         do m.each_reverse |k, v| {
 590             if n == uint::max_value - 5000 { false } else {
 591                 assert!(n > uint::max_value - 5000);
 592
 593                 assert_eq!(*k, n);
 594                 assert_eq!(*v, n / 2);
 595                 n -= 1;
 596                 true
 597             }
 598         };
 599     }
 600
 601     #[test]
 602     fn test_swap() {
 603         let mut m = TrieMap::new();
 604         assert_eq!(m.swap(1, 2), None);
 605         assert_eq!(m.swap(1, 3), Some(2));
 606         assert_eq!(m.swap(1, 4), Some(3));
 607     }
 608
 609     #[test]
 610     fn test_pop() {
 611         let mut m = TrieMap::new();
 612         m.insert(1, 2);
 613         assert_eq!(m.pop(&1), Some(2));
 614         assert_eq!(m.pop(&1), None);
 615     }
 616
 617     #[test]
 618     fn test_from_iter() {
 619         let xs = ~[(1u, 1i), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
 620
 621         let map: TrieMap<int> = xs.iter().transform(|&x| x).collect();
 622
 623         for &(k, v) in xs.iter() {
 624             assert_eq!(map.find(&k), Some(&v));
 625         }
 626     }
 627
 628     #[test]
 629     fn test_iteration() {
 630         let empty_map : TrieMap<uint> = TrieMap::new();
 631         assert_eq!(empty_map.iter().next(), None);
 632
 633         let first = uint::max_value - 10000;
 634         let last = uint::max_value;
 635
 636         let mut map = TrieMap::new();
 637         do uint::range_rev(last, first) |x| {
 638             map.insert(x, x / 2);
 639             true
 640         };
 641
 642         let mut i = 0;
 643         for (k, &v) in map.iter() {
 644             assert_eq!(k, first + i);
 645             assert_eq!(v, k / 2);
 646             i += 1;
 647         }
 648         assert_eq!(i, last - first);
 649     }
 650 }
 651
 652 #[cfg(test)]
 653 mod test_set {
 654     use super::*;
 655     use prelude::*;
 656     use uint;
 657
 658     #[test]
 659     fn test_sane_chunk() {
 660         let x = 1;
 661         let y = 1 << (uint::bits - 1);
 662
 663         let mut trie = TrieSet::new();
 664
 665         assert!(trie.insert(x));
 666         assert!(trie.insert(y));
 667
 668         assert_eq!(trie.len(), 2);
 669
 670         let expected = [x, y];
 671
 672         let mut i = 0;
 673
 674         do trie.each |x| {
 675             assert_eq!(expected[i], *x);
 676             i += 1;
 677             true
 678         };
 679     }
 680
 681     #[test]
 682     fn test_from_iter() {
 683         let xs = ~[9u, 8, 7, 6, 5, 4, 3, 2, 1];
 684
 685         let set: TrieSet = xs.iter().transform(|&x| x).collect();
 686
 687         for x in xs.iter() {
 688             assert!(set.contains(x));
 689         }
 690     }
 691 }