]> git.lizzy.rs Git - rust.git/blob - src/libcollections/smallintmap.rs
auto merge of #15945 : treeman/rust/doc-smallint-update, r=alexcrichton
[rust.git] / src / libcollections / smallintmap.rs
1 // Copyright 2012-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 //! A simple map based on a vector for small integer keys. Space requirements
12 //! are O(highest integer key).
13
14 #![allow(missing_doc)]
15
16 use core::prelude::*;
17
18 use core::default::Default;
19 use core::fmt;
20 use core::iter::{Enumerate, FilterMap};
21 use core::mem::replace;
22
23 use {Collection, Mutable, Map, MutableMap, MutableSeq};
24 use {vec, slice};
25 use vec::Vec;
26
27 /// A map optimized for small integer keys.
28 ///
29 /// # Example
30 ///
31 /// ```
32 /// use std::collections::SmallIntMap;
33 ///
34 /// let mut months = SmallIntMap::new();
35 /// months.insert(1, "Jan");
36 /// months.insert(2, "Feb");
37 /// months.insert(3, "Mar");
38 ///
39 /// if !months.contains_key(&12) {
40 ///     println!("The end is near!");
41 /// }
42 ///
43 /// assert_eq!(months.find(&1), Some(&"Jan"));
44 ///
45 /// match months.find_mut(&3) {
46 ///     Some(value) => *value = "Venus",
47 ///     None => (),
48 /// }
49 ///
50 /// assert_eq!(months.find(&3), Some(&"Venus"));
51 ///
52 /// // Print out all months
53 /// for (key, value) in months.iter() {
54 ///     println!("month {} is {}", key, value);
55 /// }
56 ///
57 /// months.clear();
58 /// assert!(months.is_empty());
59 /// ```
60 pub struct SmallIntMap<T> {
61     v: Vec<Option<T>>,
62 }
63
64 impl<V> Collection for SmallIntMap<V> {
65     /// Return the number of elements in the map.
66     fn len(&self) -> uint {
67         self.v.iter().filter(|elt| elt.is_some()).count()
68     }
69
70     /// Return `true` if there are no elements in the map.
71     fn is_empty(&self) -> bool {
72         self.v.iter().all(|elt| elt.is_none())
73     }
74 }
75
76 impl<V> Mutable for SmallIntMap<V> {
77     /// Clear the map, removing all key-value pairs.
78     fn clear(&mut self) { self.v.clear() }
79 }
80
81 impl<V> Map<uint, V> for SmallIntMap<V> {
82     /// Return a reference to the value corresponding to the key.
83     fn find<'a>(&'a self, key: &uint) -> Option<&'a V> {
84         if *key < self.v.len() {
85             match *self.v.get(*key) {
86               Some(ref value) => Some(value),
87               None => None
88             }
89         } else {
90             None
91         }
92     }
93 }
94
95 impl<V> MutableMap<uint, V> for SmallIntMap<V> {
96     /// Return a mutable reference to the value corresponding to the key.
97     fn find_mut<'a>(&'a mut self, key: &uint) -> Option<&'a mut V> {
98         if *key < self.v.len() {
99             match *self.v.get_mut(*key) {
100               Some(ref mut value) => Some(value),
101               None => None
102             }
103         } else {
104             None
105         }
106     }
107
108     /// Insert a key-value pair into the map. An existing value for a
109     /// key is replaced by the new value. Return `true` if the key did
110     /// not already exist in the map.
111     fn insert(&mut self, key: uint, value: V) -> bool {
112         let exists = self.contains_key(&key);
113         let len = self.v.len();
114         if len <= key {
115             self.v.grow_fn(key - len + 1, |_| None);
116         }
117         *self.v.get_mut(key) = Some(value);
118         !exists
119     }
120
121     /// Remove a key-value pair from the map. Return `true` if the key
122     /// was present in the map, otherwise `false`.
123     fn remove(&mut self, key: &uint) -> bool {
124         self.pop(key).is_some()
125     }
126
127     /// Insert a key-value pair from the map. If the key already had a value
128     /// present in the map, that value is returned. Otherwise `None` is returned.
129     fn swap(&mut self, key: uint, value: V) -> Option<V> {
130         match self.find_mut(&key) {
131             Some(loc) => { return Some(replace(loc, value)); }
132             None => ()
133         }
134         self.insert(key, value);
135         return None;
136     }
137
138     /// Removes a key from the map, returning the value at the key if the key
139     /// was previously in the map.
140     fn pop(&mut self, key: &uint) -> Option<V> {
141         if *key >= self.v.len() {
142             return None;
143         }
144         self.v.get_mut(*key).take()
145     }
146 }
147
148 impl<V> Default for SmallIntMap<V> {
149     #[inline]
150     fn default() -> SmallIntMap<V> { SmallIntMap::new() }
151 }
152
153 impl<V> SmallIntMap<V> {
154     /// Create an empty SmallIntMap.
155     ///
156     /// # Example
157     ///
158     /// ```
159     /// use std::collections::SmallIntMap;
160     /// let mut map: SmallIntMap<&str> = SmallIntMap::new();
161     /// ```
162     pub fn new() -> SmallIntMap<V> { SmallIntMap{v: vec!()} }
163
164     /// Create an empty SmallIntMap with space for at least `capacity` elements
165     /// before resizing.
166     ///
167     /// # Example
168     ///
169     /// ```
170     /// use std::collections::SmallIntMap;
171     /// let mut map: SmallIntMap<&str> = SmallIntMap::with_capacity(10);
172     /// ```
173     pub fn with_capacity(capacity: uint) -> SmallIntMap<V> {
174         SmallIntMap { v: Vec::with_capacity(capacity) }
175     }
176
177     /// Retrieves a value for the given key.
178     /// See [`find`](../trait.Map.html#tymethod.find) for a non-failing alternative.
179     ///
180     /// # Failure
181     ///
182     /// Fails if the key is not present.
183     ///
184     /// # Example
185     ///
186     /// ```
187     /// use std::collections::SmallIntMap;
188     ///
189     /// let mut map = SmallIntMap::new();
190     /// map.insert(1, "a");
191     /// assert_eq!(map.get(&1), &"a");
192     /// ```
193     pub fn get<'a>(&'a self, key: &uint) -> &'a V {
194         self.find(key).expect("key not present")
195     }
196
197     /// An iterator visiting all key-value pairs in ascending order by the keys.
198     /// Iterator element type is `(uint, &'r V)`.
199     ///
200     /// # Example
201     ///
202     /// ```
203     /// use std::collections::SmallIntMap;
204     ///
205     /// let mut map = SmallIntMap::new();
206     /// map.insert(1, "a");
207     /// map.insert(3, "c");
208     /// map.insert(2, "b");
209     ///
210     /// // Print `1: a` then `2: b` then `3: c`
211     /// for (key, value) in map.iter() {
212     ///     println!("{}: {}", key, value);
213     /// }
214     /// ```
215     pub fn iter<'r>(&'r self) -> Entries<'r, V> {
216         Entries {
217             front: 0,
218             back: self.v.len(),
219             iter: self.v.iter()
220         }
221     }
222
223     /// An iterator visiting all key-value pairs in ascending order by the keys,
224     /// with mutable references to the values
225     /// Iterator element type is `(uint, &'r mut V)`.
226     ///
227     /// # Example
228     ///
229     /// ```
230     /// use std::collections::SmallIntMap;
231     ///
232     /// let mut map = SmallIntMap::new();
233     /// map.insert(1, "a");
234     /// map.insert(2, "b");
235     /// map.insert(3, "c");
236     ///
237     /// for (key, value) in map.mut_iter() {
238     ///     *value = "x";
239     /// }
240     ///
241     /// for (key, value) in map.iter() {
242     ///     assert_eq!(value, &"x");
243     /// }
244     /// ```
245     pub fn mut_iter<'r>(&'r mut self) -> MutEntries<'r, V> {
246         MutEntries {
247             front: 0,
248             back: self.v.len(),
249             iter: self.v.mut_iter()
250         }
251     }
252
253     /// Empties the hash map, moving all values into the specified closure.
254     ///
255     /// # Example
256     ///
257     /// ```
258     /// use std::collections::SmallIntMap;
259     ///
260     /// let mut map = SmallIntMap::new();
261     /// map.insert(1, "a");
262     /// map.insert(3, "c");
263     /// map.insert(2, "b");
264     ///
265     /// // Not possible with .iter()
266     /// let vec: Vec<(uint, &str)> = map.move_iter().collect();
267     ///
268     /// assert_eq!(vec, vec![(1, "a"), (2, "b"), (3, "c")]);
269     /// ```
270     pub fn move_iter(&mut self)
271         -> FilterMap<(uint, Option<V>), (uint, V),
272                 Enumerate<vec::MoveItems<Option<V>>>>
273     {
274         let values = replace(&mut self.v, vec!());
275         values.move_iter().enumerate().filter_map(|(i, v)| {
276             v.map(|v| (i, v))
277         })
278     }
279 }
280
281 impl<V:Clone> SmallIntMap<V> {
282     /// Update a value in the map. If the key already exists in the map,
283     /// modify the value with `ff` taking `oldval, newval`.
284     /// Otherwise set the value to `newval`.
285     /// Return `true` if the key did not already exist in the map.
286     ///
287     /// # Example
288     ///
289     /// ```
290     /// use std::collections::SmallIntMap;
291     ///
292     /// let mut map = SmallIntMap::new();
293     ///
294     /// // Key does not exist, will do a simple insert
295     /// assert!(map.update(1, vec![1i, 2], |old, new| old.append(new.as_slice())));
296     /// assert_eq!(map.get(&1), &vec![1i, 2]);
297     ///
298     /// // Key exists, update the value
299     /// assert!(!map.update(1, vec![3i, 4], |old, new| old.append(new.as_slice())));
300     /// assert_eq!(map.get(&1), &vec![1i, 2, 3, 4]);
301     /// ```
302     pub fn update(&mut self, key: uint, newval: V, ff: |V, V| -> V) -> bool {
303         self.update_with_key(key, newval, |_k, v, v1| ff(v,v1))
304     }
305
306     /// Update a value in the map. If the key already exists in the map,
307     /// modify the value with `ff` taking `key, oldval, newval`.
308     /// Otherwise set the value to `newval`.
309     /// Return `true` if the key did not already exist in the map.
310     ///
311     /// # Example
312     ///
313     /// ```
314     /// use std::collections::SmallIntMap;
315     ///
316     /// let mut map = SmallIntMap::new();
317     ///
318     /// // Key does not exist, will do a simple insert
319     /// assert!(map.update_with_key(7, 10, |key, old, new| (old + new) % key));
320     /// assert_eq!(map.get(&7), &10);
321     ///
322     /// // Key exists, update the value
323     /// assert!(!map.update_with_key(7, 20, |key, old, new| (old + new) % key));
324     /// assert_eq!(map.get(&7), &2);
325     /// ```
326     pub fn update_with_key(&mut self,
327                            key: uint,
328                            val: V,
329                            ff: |uint, V, V| -> V)
330                            -> bool {
331         let new_val = match self.find(&key) {
332             None => val,
333             Some(orig) => ff(key, (*orig).clone(), val)
334         };
335         self.insert(key, new_val)
336     }
337 }
338
339 impl<V: fmt::Show> fmt::Show for SmallIntMap<V> {
340     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
341         try!(write!(f, "{{"));
342
343         for (i, (k, v)) in self.iter().enumerate() {
344             if i != 0 { try!(write!(f, ", ")); }
345             try!(write!(f, "{}: {}", k, *v));
346         }
347
348         write!(f, "}}")
349     }
350 }
351
352 macro_rules! iterator {
353     (impl $name:ident -> $elem:ty, $getter:ident) => {
354         impl<'a, T> Iterator<$elem> for $name<'a, T> {
355             #[inline]
356             fn next(&mut self) -> Option<$elem> {
357                 while self.front < self.back {
358                     match self.iter.next() {
359                         Some(elem) => {
360                             if elem.is_some() {
361                                 let index = self.front;
362                                 self.front += 1;
363                                 return Some((index, elem. $getter ()));
364                             }
365                         }
366                         _ => ()
367                     }
368                     self.front += 1;
369                 }
370                 None
371             }
372
373             #[inline]
374             fn size_hint(&self) -> (uint, Option<uint>) {
375                 (0, Some(self.back - self.front))
376             }
377         }
378     }
379 }
380
381 macro_rules! double_ended_iterator {
382     (impl $name:ident -> $elem:ty, $getter:ident) => {
383         impl<'a, T> DoubleEndedIterator<$elem> for $name<'a, T> {
384             #[inline]
385             fn next_back(&mut self) -> Option<$elem> {
386                 while self.front < self.back {
387                     match self.iter.next_back() {
388                         Some(elem) => {
389                             if elem.is_some() {
390                                 self.back -= 1;
391                                 return Some((self.back, elem. $getter ()));
392                             }
393                         }
394                         _ => ()
395                     }
396                     self.back -= 1;
397                 }
398                 None
399             }
400         }
401     }
402 }
403
404 /// Forward iterator over a map.
405 pub struct Entries<'a, T> {
406     front: uint,
407     back: uint,
408     iter: slice::Items<'a, Option<T>>
409 }
410
411 iterator!(impl Entries -> (uint, &'a T), get_ref)
412 double_ended_iterator!(impl Entries -> (uint, &'a T), get_ref)
413
414 /// Forward iterator over the key-value pairs of a map, with the
415 /// values being mutable.
416 pub struct MutEntries<'a, T> {
417     front: uint,
418     back: uint,
419     iter: slice::MutItems<'a, Option<T>>
420 }
421
422 iterator!(impl MutEntries -> (uint, &'a mut T), get_mut_ref)
423 double_ended_iterator!(impl MutEntries -> (uint, &'a mut T), get_mut_ref)
424
425 #[cfg(test)]
426 mod test_map {
427     use std::prelude::*;
428
429     use {Map, MutableMap, Mutable};
430     use super::SmallIntMap;
431
432     #[test]
433     fn test_find_mut() {
434         let mut m = SmallIntMap::new();
435         assert!(m.insert(1, 12i));
436         assert!(m.insert(2, 8));
437         assert!(m.insert(5, 14));
438         let new = 100;
439         match m.find_mut(&5) {
440             None => fail!(), Some(x) => *x = new
441         }
442         assert_eq!(m.find(&5), Some(&new));
443     }
444
445     #[test]
446     fn test_len() {
447         let mut map = SmallIntMap::new();
448         assert_eq!(map.len(), 0);
449         assert!(map.is_empty());
450         assert!(map.insert(5, 20i));
451         assert_eq!(map.len(), 1);
452         assert!(!map.is_empty());
453         assert!(map.insert(11, 12));
454         assert_eq!(map.len(), 2);
455         assert!(!map.is_empty());
456         assert!(map.insert(14, 22));
457         assert_eq!(map.len(), 3);
458         assert!(!map.is_empty());
459     }
460
461     #[test]
462     fn test_clear() {
463         let mut map = SmallIntMap::new();
464         assert!(map.insert(5, 20i));
465         assert!(map.insert(11, 12));
466         assert!(map.insert(14, 22));
467         map.clear();
468         assert!(map.is_empty());
469         assert!(map.find(&5).is_none());
470         assert!(map.find(&11).is_none());
471         assert!(map.find(&14).is_none());
472     }
473
474     #[test]
475     fn test_insert_with_key() {
476         let mut map = SmallIntMap::new();
477
478         // given a new key, initialize it with this new count,
479         // given an existing key, add more to its count
480         fn add_more_to_count(_k: uint, v0: uint, v1: uint) -> uint {
481             v0 + v1
482         }
483
484         fn add_more_to_count_simple(v0: uint, v1: uint) -> uint {
485             v0 + v1
486         }
487
488         // count integers
489         map.update(3, 1, add_more_to_count_simple);
490         map.update_with_key(9, 1, add_more_to_count);
491         map.update(3, 7, add_more_to_count_simple);
492         map.update_with_key(5, 3, add_more_to_count);
493         map.update_with_key(3, 2, add_more_to_count);
494
495         // check the total counts
496         assert_eq!(map.find(&3).unwrap(), &10);
497         assert_eq!(map.find(&5).unwrap(), &3);
498         assert_eq!(map.find(&9).unwrap(), &1);
499
500         // sadly, no sevens were counted
501         assert!(map.find(&7).is_none());
502     }
503
504     #[test]
505     fn test_swap() {
506         let mut m = SmallIntMap::new();
507         assert_eq!(m.swap(1, 2i), None);
508         assert_eq!(m.swap(1, 3i), Some(2));
509         assert_eq!(m.swap(1, 4i), Some(3));
510     }
511
512     #[test]
513     fn test_pop() {
514         let mut m = SmallIntMap::new();
515         m.insert(1, 2i);
516         assert_eq!(m.pop(&1), Some(2));
517         assert_eq!(m.pop(&1), None);
518     }
519
520     #[test]
521     fn test_iterator() {
522         let mut m = SmallIntMap::new();
523
524         assert!(m.insert(0, 1i));
525         assert!(m.insert(1, 2));
526         assert!(m.insert(3, 5));
527         assert!(m.insert(6, 10));
528         assert!(m.insert(10, 11));
529
530         let mut it = m.iter();
531         assert_eq!(it.size_hint(), (0, Some(11)));
532         assert_eq!(it.next().unwrap(), (0, &1));
533         assert_eq!(it.size_hint(), (0, Some(10)));
534         assert_eq!(it.next().unwrap(), (1, &2));
535         assert_eq!(it.size_hint(), (0, Some(9)));
536         assert_eq!(it.next().unwrap(), (3, &5));
537         assert_eq!(it.size_hint(), (0, Some(7)));
538         assert_eq!(it.next().unwrap(), (6, &10));
539         assert_eq!(it.size_hint(), (0, Some(4)));
540         assert_eq!(it.next().unwrap(), (10, &11));
541         assert_eq!(it.size_hint(), (0, Some(0)));
542         assert!(it.next().is_none());
543     }
544
545     #[test]
546     fn test_iterator_size_hints() {
547         let mut m = SmallIntMap::new();
548
549         assert!(m.insert(0, 1i));
550         assert!(m.insert(1, 2));
551         assert!(m.insert(3, 5));
552         assert!(m.insert(6, 10));
553         assert!(m.insert(10, 11));
554
555         assert_eq!(m.iter().size_hint(), (0, Some(11)));
556         assert_eq!(m.iter().rev().size_hint(), (0, Some(11)));
557         assert_eq!(m.mut_iter().size_hint(), (0, Some(11)));
558         assert_eq!(m.mut_iter().rev().size_hint(), (0, Some(11)));
559     }
560
561     #[test]
562     fn test_mut_iterator() {
563         let mut m = SmallIntMap::new();
564
565         assert!(m.insert(0, 1i));
566         assert!(m.insert(1, 2));
567         assert!(m.insert(3, 5));
568         assert!(m.insert(6, 10));
569         assert!(m.insert(10, 11));
570
571         for (k, v) in m.mut_iter() {
572             *v += k as int;
573         }
574
575         let mut it = m.iter();
576         assert_eq!(it.next().unwrap(), (0, &1));
577         assert_eq!(it.next().unwrap(), (1, &3));
578         assert_eq!(it.next().unwrap(), (3, &8));
579         assert_eq!(it.next().unwrap(), (6, &16));
580         assert_eq!(it.next().unwrap(), (10, &21));
581         assert!(it.next().is_none());
582     }
583
584     #[test]
585     fn test_rev_iterator() {
586         let mut m = SmallIntMap::new();
587
588         assert!(m.insert(0, 1i));
589         assert!(m.insert(1, 2));
590         assert!(m.insert(3, 5));
591         assert!(m.insert(6, 10));
592         assert!(m.insert(10, 11));
593
594         let mut it = m.iter().rev();
595         assert_eq!(it.next().unwrap(), (10, &11));
596         assert_eq!(it.next().unwrap(), (6, &10));
597         assert_eq!(it.next().unwrap(), (3, &5));
598         assert_eq!(it.next().unwrap(), (1, &2));
599         assert_eq!(it.next().unwrap(), (0, &1));
600         assert!(it.next().is_none());
601     }
602
603     #[test]
604     fn test_mut_rev_iterator() {
605         let mut m = SmallIntMap::new();
606
607         assert!(m.insert(0, 1i));
608         assert!(m.insert(1, 2));
609         assert!(m.insert(3, 5));
610         assert!(m.insert(6, 10));
611         assert!(m.insert(10, 11));
612
613         for (k, v) in m.mut_iter().rev() {
614             *v += k as int;
615         }
616
617         let mut it = m.iter();
618         assert_eq!(it.next().unwrap(), (0, &1));
619         assert_eq!(it.next().unwrap(), (1, &3));
620         assert_eq!(it.next().unwrap(), (3, &8));
621         assert_eq!(it.next().unwrap(), (6, &16));
622         assert_eq!(it.next().unwrap(), (10, &21));
623         assert!(it.next().is_none());
624     }
625
626     #[test]
627     fn test_move_iter() {
628         let mut m = SmallIntMap::new();
629         m.insert(1, box 2i);
630         let mut called = false;
631         for (k, v) in m.move_iter() {
632             assert!(!called);
633             called = true;
634             assert_eq!(k, 1);
635             assert_eq!(v, box 2i);
636         }
637         assert!(called);
638         m.insert(2, box 1i);
639     }
640
641     #[test]
642     fn test_show() {
643         let mut map = SmallIntMap::new();
644         let empty = SmallIntMap::<int>::new();
645
646         map.insert(1, 2i);
647         map.insert(3, 4i);
648
649         let map_str = map.to_string();
650         let map_str = map_str.as_slice();
651         assert!(map_str == "{1: 2, 3: 4}" || map_str == "{3: 4, 1: 2}");
652         assert_eq!(format!("{}", empty), "{}".to_string());
653     }
654 }
655
656 #[cfg(test)]
657 mod bench {
658     extern crate test;
659     use self::test::Bencher;
660     use super::SmallIntMap;
661     use deque::bench::{insert_rand_n, insert_seq_n, find_rand_n, find_seq_n};
662
663     // Find seq
664     #[bench]
665     pub fn insert_rand_100(b: &mut Bencher) {
666         let mut m : SmallIntMap<uint> = SmallIntMap::new();
667         insert_rand_n(100, &mut m, b);
668     }
669
670     #[bench]
671     pub fn insert_rand_10_000(b: &mut Bencher) {
672         let mut m : SmallIntMap<uint> = SmallIntMap::new();
673         insert_rand_n(10_000, &mut m, b);
674     }
675
676     // Insert seq
677     #[bench]
678     pub fn insert_seq_100(b: &mut Bencher) {
679         let mut m : SmallIntMap<uint> = SmallIntMap::new();
680         insert_seq_n(100, &mut m, b);
681     }
682
683     #[bench]
684     pub fn insert_seq_10_000(b: &mut Bencher) {
685         let mut m : SmallIntMap<uint> = SmallIntMap::new();
686         insert_seq_n(10_000, &mut m, b);
687     }
688
689     // Find rand
690     #[bench]
691     pub fn find_rand_100(b: &mut Bencher) {
692         let mut m : SmallIntMap<uint> = SmallIntMap::new();
693         find_rand_n(100, &mut m, b);
694     }
695
696     #[bench]
697     pub fn find_rand_10_000(b: &mut Bencher) {
698         let mut m : SmallIntMap<uint> = SmallIntMap::new();
699         find_rand_n(10_000, &mut m, b);
700     }
701
702     // Find seq
703     #[bench]
704     pub fn find_seq_100(b: &mut Bencher) {
705         let mut m : SmallIntMap<uint> = SmallIntMap::new();
706         find_seq_n(100, &mut m, b);
707     }
708
709     #[bench]
710     pub fn find_seq_10_000(b: &mut Bencher) {
711         let mut m : SmallIntMap<uint> = SmallIntMap::new();
712         find_seq_n(10_000, &mut m, b);
713     }
714 }