src/libextra/smallintmap.rs

   1 // Copyright 2012 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 /*!
  12  * A simple map based on a vector for small integer keys. Space requirements
  13  * are O(highest integer key).
  14  */
  15
  16 #[allow(missing_doc)];
  17
  18 use std::iterator::{Iterator, IteratorUtil, Enumerate, FilterMap, Invert};
  19 use std::uint;
  20 use std::util::replace;
  21 use std::vec::{VecIterator, VecMutIterator};
  22 use std::vec;
  23
  24 #[allow(missing_doc)]
  25 pub struct SmallIntMap<T> {
  26     priv v: ~[Option<T>],
  27 }
  28
  29 impl<V> Container for SmallIntMap<V> {
  30     /// Return the number of elements in the map
  31     fn len(&self) -> uint {
  32         let mut sz = 0;
  33         for i in range(0u, self.v.len()) {
  34             match self.v[i] {
  35                 Some(_) => sz += 1,
  36                 None => {}
  37             }
  38         }
  39         sz
  40     }
  41 }
  42
  43 impl<V> Mutable for SmallIntMap<V> {
  44     /// Clear the map, removing all key-value pairs.
  45     fn clear(&mut self) { self.v.clear() }
  46 }
  47
  48 impl<V> Map<uint, V> for SmallIntMap<V> {
  49     /// Return a reference to the value corresponding to the key
  50     fn find<'a>(&'a self, key: &uint) -> Option<&'a V> {
  51         if *key < self.v.len() {
  52             match self.v[*key] {
  53               Some(ref value) => Some(value),
  54               None => None
  55             }
  56         } else {
  57             None
  58         }
  59     }
  60 }
  61
  62 impl<V> MutableMap<uint, V> for SmallIntMap<V> {
  63     /// Return a mutable reference to the value corresponding to the key
  64     fn find_mut<'a>(&'a mut self, key: &uint) -> Option<&'a mut V> {
  65         if *key < self.v.len() {
  66             match self.v[*key] {
  67               Some(ref mut value) => Some(value),
  68               None => None
  69             }
  70         } else {
  71             None
  72         }
  73     }
  74
  75     /// Insert a key-value pair into the map. An existing value for a
  76     /// key is replaced by the new value. Return true if the key did
  77     /// not already exist in the map.
  78     fn insert(&mut self, key: uint, value: V) -> bool {
  79         let exists = self.contains_key(&key);
  80         let len = self.v.len();
  81         if len <= key {
  82             self.v.grow_fn(key - len + 1, |_| None);
  83         }
  84         self.v[key] = Some(value);
  85         !exists
  86     }
  87
  88     /// Remove a key-value pair from the map. Return true if the key
  89     /// was present in the map, otherwise false.
  90     fn remove(&mut self, key: &uint) -> bool {
  91         self.pop(key).is_some()
  92     }
  93
  94     /// Insert a key-value pair from the map. If the key already had a value
  95     /// present in the map, that value is returned. Otherwise None is returned.
  96     fn swap(&mut self, key: uint, value: V) -> Option<V> {
  97         match self.find_mut(&key) {
  98             Some(loc) => { return Some(replace(loc, value)); }
  99             None => ()
 100         }
 101         self.insert(key, value);
 102         return None;
 103     }
 104
 105     /// Removes a key from the map, returning the value at the key if the key
 106     /// was previously in the map.
 107     fn pop(&mut self, key: &uint) -> Option<V> {
 108         if *key >= self.v.len() {
 109             return None;
 110         }
 111         self.v[*key].take()
 112     }
 113 }
 114
 115 impl<V> SmallIntMap<V> {
 116     /// Create an empty SmallIntMap
 117     pub fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[]} }
 118
 119     /// Visit all key-value pairs in order
 120     pub fn each<'a>(&'a self, it: &fn(&uint, &'a V) -> bool) -> bool {
 121         for i in range(0u, self.v.len()) {
 122             match self.v[i] {
 123               Some(ref elt) => if !it(&i, elt) { return false; },
 124               None => ()
 125             }
 126         }
 127         true
 128     }
 129
 130     /// Visit all keys in order
 131     pub fn each_key(&self, blk: &fn(key: &uint) -> bool) -> bool {
 132         self.each(|k, _| blk(k))
 133     }
 134
 135     /// Visit all values in order
 136     pub fn each_value<'a>(&'a self, blk: &fn(value: &'a V) -> bool) -> bool {
 137         self.each(|_, v| blk(v))
 138     }
 139
 140     /// Iterate over the map and mutate the contained values
 141     pub fn mutate_values(&mut self, it: &fn(&uint, &mut V) -> bool) -> bool {
 142         for i in range(0, self.v.len()) {
 143             match self.v[i] {
 144               Some(ref mut elt) => if !it(&i, elt) { return false; },
 145               None => ()
 146             }
 147         }
 148         true
 149     }
 150
 151     /// Visit all key-value pairs in reverse order
 152     pub fn each_reverse<'a>(&'a self, it: &fn(uint, &'a V) -> bool) -> bool {
 153         do uint::range_rev(self.v.len(), 0) |i| {
 154             match self.v[i] {
 155               Some(ref elt) => it(i, elt),
 156               None => true
 157             }
 158         }
 159     }
 160
 161     pub fn get<'a>(&'a self, key: &uint) -> &'a V {
 162         self.find(key).expect("key not present")
 163     }
 164
 165     /// An iterator visiting all key-value pairs in ascending order by the keys.
 166     /// Iterator element type is (uint, &'r V)
 167     pub fn iter<'r>(&'r self) -> SmallIntMapIterator<'r, V> {
 168         SmallIntMapIterator {
 169             front: 0,
 170             back: self.v.len(),
 171             iter: self.v.iter()
 172         }
 173     }
 174
 175     /// An iterator visiting all key-value pairs in ascending order by the keys,
 176     /// with mutable references to the values
 177     /// Iterator element type is (uint, &'r mut V)
 178     pub fn mut_iter<'r>(&'r mut self) -> SmallIntMapMutIterator<'r, V> {
 179         SmallIntMapMutIterator {
 180             front: 0,
 181             back: self.v.len(),
 182             iter: self.v.mut_iter()
 183         }
 184     }
 185
 186     /// An iterator visiting all key-value pairs in descending order by the keys.
 187     /// Iterator element type is (uint, &'r V)
 188     pub fn rev_iter<'r>(&'r self) -> SmallIntMapRevIterator<'r, V> {
 189         self.iter().invert()
 190     }
 191
 192     /// An iterator visiting all key-value pairs in descending order by the keys,
 193     /// with mutable references to the values
 194     /// Iterator element type is (uint, &'r mut V)
 195     pub fn mut_rev_iter<'r>(&'r mut self) -> SmallIntMapMutRevIterator <'r, V> {
 196         self.mut_iter().invert()
 197     }
 198
 199     /// Empties the hash map, moving all values into the specified closure
 200     pub fn consume(&mut self)
 201         -> FilterMap<(uint, Option<V>), (uint, V),
 202                 Enumerate<vec::ConsumeIterator<Option<V>>>>
 203     {
 204         let values = replace(&mut self.v, ~[]);
 205         values.consume_iter().enumerate().filter_map(|(i, v)| {
 206             v.map_move(|v| (i, v))
 207         })
 208     }
 209 }
 210
 211 impl<V:Clone> SmallIntMap<V> {
 212     pub fn update_with_key(&mut self, key: uint, val: V,
 213                            ff: &fn(uint, V, V) -> V) -> bool {
 214         let new_val = match self.find(&key) {
 215             None => val,
 216             Some(orig) => ff(key, (*orig).clone(), val)
 217         };
 218         self.insert(key, new_val)
 219     }
 220
 221     pub fn update(&mut self, key: uint, newval: V, ff: &fn(V, V) -> V)
 222                   -> bool {
 223         self.update_with_key(key, newval, |_k, v, v1| ff(v,v1))
 224     }
 225 }
 226
 227
 228 macro_rules! iterator {
 229     (impl $name:ident -> $elem:ty, $getter:ident) => {
 230         impl<'self, T> Iterator<$elem> for $name<'self, T> {
 231             #[inline]
 232             fn next(&mut self) -> Option<$elem> {
 233                 while self.front < self.back {
 234                     match self.iter.next() {
 235                         Some(elem) => {
 236                             if elem.is_some() {
 237                                 let index = self.front;
 238                                 self.front += 1;
 239                                 return Some((index, elem. $getter ()));
 240                             }
 241                         }
 242                         _ => ()
 243                     }
 244                     self.front += 1;
 245                 }
 246                 None
 247             }
 248
 249             #[inline]
 250             fn size_hint(&self) -> (uint, Option<uint>) {
 251                 (0, Some(self.back - self.front))
 252             }
 253         }
 254     }
 255 }
 256
 257 macro_rules! double_ended_iterator {
 258     (impl $name:ident -> $elem:ty, $getter:ident) => {
 259         impl<'self, T> DoubleEndedIterator<$elem> for $name<'self, T> {
 260             #[inline]
 261             fn next_back(&mut self) -> Option<$elem> {
 262                 while self.front < self.back {
 263                     match self.iter.next_back() {
 264                         Some(elem) => {
 265                             if elem.is_some() {
 266                                 self.back -= 1;
 267                                 return Some((self.back, elem. $getter ()));
 268                             }
 269                         }
 270                         _ => ()
 271                     }
 272                     self.back -= 1;
 273                 }
 274                 None
 275             }
 276         }
 277     }
 278 }
 279
 280 pub struct SmallIntMapIterator<'self, T> {
 281     priv front: uint,
 282     priv back: uint,
 283     priv iter: VecIterator<'self, Option<T>>
 284 }
 285
 286 iterator!(impl SmallIntMapIterator -> (uint, &'self T), get_ref)
 287 double_ended_iterator!(impl SmallIntMapIterator -> (uint, &'self T), get_ref)
 288 pub type SmallIntMapRevIterator<'self, T> = Invert<SmallIntMapIterator<'self, T>>;
 289
 290 pub struct SmallIntMapMutIterator<'self, T> {
 291     priv front: uint,
 292     priv back: uint,
 293     priv iter: VecMutIterator<'self, Option<T>>
 294 }
 295
 296 iterator!(impl SmallIntMapMutIterator -> (uint, &'self mut T), get_mut_ref)
 297 double_ended_iterator!(impl SmallIntMapMutIterator -> (uint, &'self mut T), get_mut_ref)
 298 pub type SmallIntMapMutRevIterator<'self, T> = Invert<SmallIntMapMutIterator<'self, T>>;
 299
 300 #[cfg(test)]
 301 mod test_map {
 302
 303     use super::SmallIntMap;
 304
 305     #[test]
 306     fn test_find_mut() {
 307         let mut m = SmallIntMap::new();
 308         assert!(m.insert(1, 12));
 309         assert!(m.insert(2, 8));
 310         assert!(m.insert(5, 14));
 311         let new = 100;
 312         match m.find_mut(&5) {
 313             None => fail!(), Some(x) => *x = new
 314         }
 315         assert_eq!(m.find(&5), Some(&new));
 316     }
 317
 318     #[test]
 319     fn test_len() {
 320         let mut map = SmallIntMap::new();
 321         assert_eq!(map.len(), 0);
 322         assert!(map.is_empty());
 323         assert!(map.insert(5, 20));
 324         assert_eq!(map.len(), 1);
 325         assert!(!map.is_empty());
 326         assert!(map.insert(11, 12));
 327         assert_eq!(map.len(), 2);
 328         assert!(!map.is_empty());
 329         assert!(map.insert(14, 22));
 330         assert_eq!(map.len(), 3);
 331         assert!(!map.is_empty());
 332     }
 333
 334     #[test]
 335     fn test_clear() {
 336         let mut map = SmallIntMap::new();
 337         assert!(map.insert(5, 20));
 338         assert!(map.insert(11, 12));
 339         assert!(map.insert(14, 22));
 340         map.clear();
 341         assert!(map.is_empty());
 342         assert!(map.find(&5).is_none());
 343         assert!(map.find(&11).is_none());
 344         assert!(map.find(&14).is_none());
 345     }
 346
 347     #[test]
 348     fn test_insert_with_key() {
 349         let mut map = SmallIntMap::new();
 350
 351         // given a new key, initialize it with this new count, given
 352         // given an existing key, add more to its count
 353         fn addMoreToCount(_k: uint, v0: uint, v1: uint) -> uint {
 354             v0 + v1
 355         }
 356
 357         fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
 358             v0 + v1
 359         }
 360
 361         // count integers
 362         map.update(3, 1, addMoreToCount_simple);
 363         map.update_with_key(9, 1, addMoreToCount);
 364         map.update(3, 7, addMoreToCount_simple);
 365         map.update_with_key(5, 3, addMoreToCount);
 366         map.update_with_key(3, 2, addMoreToCount);
 367
 368         // check the total counts
 369         assert_eq!(map.find(&3).unwrap(), &10);
 370         assert_eq!(map.find(&5).unwrap(), &3);
 371         assert_eq!(map.find(&9).unwrap(), &1);
 372
 373         // sadly, no sevens were counted
 374         assert!(map.find(&7).is_none());
 375     }
 376
 377     #[test]
 378     fn test_swap() {
 379         let mut m = SmallIntMap::new();
 380         assert_eq!(m.swap(1, 2), None);
 381         assert_eq!(m.swap(1, 3), Some(2));
 382         assert_eq!(m.swap(1, 4), Some(3));
 383     }
 384
 385     #[test]
 386     fn test_pop() {
 387         let mut m = SmallIntMap::new();
 388         m.insert(1, 2);
 389         assert_eq!(m.pop(&1), Some(2));
 390         assert_eq!(m.pop(&1), None);
 391     }
 392
 393     #[test]
 394     fn test_iterator() {
 395         let mut m = SmallIntMap::new();
 396
 397         assert!(m.insert(0, 1));
 398         assert!(m.insert(1, 2));
 399         assert!(m.insert(3, 5));
 400         assert!(m.insert(6, 10));
 401         assert!(m.insert(10, 11));
 402
 403         let mut it = m.iter();
 404         assert_eq!(it.size_hint(), (0, Some(11)));
 405         assert_eq!(it.next().unwrap(), (0, &1));
 406         assert_eq!(it.size_hint(), (0, Some(10)));
 407         assert_eq!(it.next().unwrap(), (1, &2));
 408         assert_eq!(it.size_hint(), (0, Some(9)));
 409         assert_eq!(it.next().unwrap(), (3, &5));
 410         assert_eq!(it.size_hint(), (0, Some(7)));
 411         assert_eq!(it.next().unwrap(), (6, &10));
 412         assert_eq!(it.size_hint(), (0, Some(4)));
 413         assert_eq!(it.next().unwrap(), (10, &11));
 414         assert_eq!(it.size_hint(), (0, Some(0)));
 415         assert!(it.next().is_none());
 416     }
 417
 418     #[test]
 419     fn test_iterator_size_hints() {
 420         let mut m = SmallIntMap::new();
 421
 422         assert!(m.insert(0, 1));
 423         assert!(m.insert(1, 2));
 424         assert!(m.insert(3, 5));
 425         assert!(m.insert(6, 10));
 426         assert!(m.insert(10, 11));
 427
 428         assert_eq!(m.iter().size_hint(), (0, Some(11)));
 429         assert_eq!(m.rev_iter().size_hint(), (0, Some(11)));
 430         assert_eq!(m.mut_iter().size_hint(), (0, Some(11)));
 431         assert_eq!(m.mut_rev_iter().size_hint(), (0, Some(11)));
 432     }
 433
 434     #[test]
 435     fn test_mut_iterator() {
 436         let mut m = SmallIntMap::new();
 437
 438         assert!(m.insert(0, 1));
 439         assert!(m.insert(1, 2));
 440         assert!(m.insert(3, 5));
 441         assert!(m.insert(6, 10));
 442         assert!(m.insert(10, 11));
 443
 444         for (k, v) in m.mut_iter() {
 445             *v += k as int;
 446         }
 447
 448         let mut it = m.iter();
 449         assert_eq!(it.next().unwrap(), (0, &1));
 450         assert_eq!(it.next().unwrap(), (1, &3));
 451         assert_eq!(it.next().unwrap(), (3, &8));
 452         assert_eq!(it.next().unwrap(), (6, &16));
 453         assert_eq!(it.next().unwrap(), (10, &21));
 454         assert!(it.next().is_none());
 455     }
 456
 457     #[test]
 458     fn test_rev_iterator() {
 459         let mut m = SmallIntMap::new();
 460
 461         assert!(m.insert(0, 1));
 462         assert!(m.insert(1, 2));
 463         assert!(m.insert(3, 5));
 464         assert!(m.insert(6, 10));
 465         assert!(m.insert(10, 11));
 466
 467         let mut it = m.rev_iter();
 468         assert_eq!(it.next().unwrap(), (10, &11));
 469         assert_eq!(it.next().unwrap(), (6, &10));
 470         assert_eq!(it.next().unwrap(), (3, &5));
 471         assert_eq!(it.next().unwrap(), (1, &2));
 472         assert_eq!(it.next().unwrap(), (0, &1));
 473         assert!(it.next().is_none());
 474     }
 475
 476     #[test]
 477     fn test_mut_rev_iterator() {
 478         let mut m = SmallIntMap::new();
 479
 480         assert!(m.insert(0, 1));
 481         assert!(m.insert(1, 2));
 482         assert!(m.insert(3, 5));
 483         assert!(m.insert(6, 10));
 484         assert!(m.insert(10, 11));
 485
 486         for (k, v) in m.mut_rev_iter() {
 487             *v += k as int;
 488         }
 489
 490         let mut it = m.iter();
 491         assert_eq!(it.next().unwrap(), (0, &1));
 492         assert_eq!(it.next().unwrap(), (1, &3));
 493         assert_eq!(it.next().unwrap(), (3, &8));
 494         assert_eq!(it.next().unwrap(), (6, &16));
 495         assert_eq!(it.next().unwrap(), (10, &21));
 496         assert!(it.next().is_none());
 497     }
 498
 499     #[test]
 500     fn test_consume() {
 501         let mut m = SmallIntMap::new();
 502         m.insert(1, ~2);
 503         let mut called = false;
 504         for (k, v) in m.consume() {
 505             assert!(!called);
 506             called = true;
 507             assert_eq!(k, 1);
 508             assert_eq!(v, ~2);
 509         }
 510         assert!(called);
 511         m.insert(2, ~1);
 512     }
 513 }
 514
 515 #[cfg(test)]
 516 mod bench {
 517
 518     use super::*;
 519     use test::BenchHarness;
 520     use container::bench::*;
 521
 522     // Find seq
 523     #[bench]
 524     pub fn insert_rand_100(bh: &mut BenchHarness) {
 525         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 526         insert_rand_n(100, &mut m, bh);
 527     }
 528
 529     #[bench]
 530     pub fn insert_rand_10_000(bh: &mut BenchHarness) {
 531         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 532         insert_rand_n(10_000, &mut m, bh);
 533     }
 534
 535     // Insert seq
 536     #[bench]
 537     pub fn insert_seq_100(bh: &mut BenchHarness) {
 538         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 539         insert_seq_n(100, &mut m, bh);
 540     }
 541
 542     #[bench]
 543     pub fn insert_seq_10_000(bh: &mut BenchHarness) {
 544         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 545         insert_seq_n(10_000, &mut m, bh);
 546     }
 547
 548     // Find rand
 549     #[bench]
 550     pub fn find_rand_100(bh: &mut BenchHarness) {
 551         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 552         find_rand_n(100, &mut m, bh);
 553     }
 554
 555     #[bench]
 556     pub fn find_rand_10_000(bh: &mut BenchHarness) {
 557         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 558         find_rand_n(10_000, &mut m, bh);
 559     }
 560
 561     // Find seq
 562     #[bench]
 563     pub fn find_seq_100(bh: &mut BenchHarness) {
 564         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 565         find_seq_n(100, &mut m, bh);
 566     }
 567
 568     #[bench]
 569     pub fn find_seq_10_000(bh: &mut BenchHarness) {
 570         let mut m : SmallIntMap<uint> = SmallIntMap::new();
 571         find_seq_n(10_000, &mut m, bh);
 572     }
 573 }