A trie in Rust
1// Copyright 2013-2026 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 based on a trie.
12
13use crate::Chunk;
14use crate::chunk::{MAX_DEPTH, SIZE};
15
16pub use self::Entry::*;
17use self::TrieNode::*;
18
19use std::borrow::Borrow;
20use std::cmp::Ordering;
21use std::fmt::{self, Debug};
22use std::hash::{Hash, Hasher};
23use std::iter;
24use std::marker::PhantomData;
25use std::mem;
26use std::ops;
27use std::ptr;
28use std::slice;
29
30/// A map implemented as a radix trie.
31///
32/// Keys are split into sequences of 4 bits, which are used to place elements in
33/// 16-entry arrays which are nested to form a tree structure. Inserted elements are placed
34/// as close to the top of the tree as possible. The most significant bits of the key are used to
35/// assign the key to a node/bucket in the first layer. If there are no other elements keyed by
36/// the same 4 bits in the first layer, a leaf node will be created in the first layer.
37/// When keys coincide, the next 4 bits are used to assign the node to a bucket in the next layer,
38/// with this process continuing until an empty spot is found or there are no more bits left in the
39/// key. As a result, the maximum depth using 32-bit `usize` keys is 8. The worst collisions occur
40/// for very small numbers. For example, 1 and 2 are identical in all but their least significant
41/// 4 bits. If both numbers are used as keys, a chain of maximum length will be created to
42/// differentiate them.
43///
44/// # Examples
45///
46/// ```
47/// let mut map = trie::Map::new();
48/// map.insert(27, "Olaf");
49/// map.insert(1, "Edgar");
50/// map.insert(13, "Ruth");
51/// map.insert(1, "Martin");
52///
53/// assert_eq!(map.len(), 3);
54/// assert_eq!(map.get(&1), Some(&"Martin"));
55///
56/// if !map.contains_key(&90) {
57/// println!("Nobody is keyed 90");
58/// }
59///
60/// // Update a key
61/// match map.get_mut(&1) {
62/// Some(value) => *value = "Olga",
63/// None => (),
64/// }
65///
66/// map.remove(&13);
67/// assert_eq!(map.len(), 2);
68///
69/// // Print the key value pairs, ordered by key.
70/// for (key, value) in map.iter() {
71/// // Prints `1: Olga` then `27: Olaf`
72/// println!("{}: {}", key, value);
73/// }
74///
75/// map.clear();
76/// assert!(map.is_empty());
77/// ```
78#[derive(Clone)]
79pub struct Map<K, V> {
80 root: InternalNode<K, V>,
81 length: usize,
82}
83
84// An internal node holds SIZE child nodes, which may themselves contain more internal nodes.
85//
86// Throughout this implementation, "idx" is used to refer to a section of key that is used
87// to access a node. The layer of the tree directly below the root corresponds to idx 0.
88#[derive(Clone)]
89struct InternalNode<K, V> {
90 // The number of direct children which are external (i.e. that store a value).
91 count: usize,
92 children: [TrieNode<K, V>; SIZE],
93}
94
95// Each child of an InternalNode may be internal, in which case nesting continues,
96// external (containing a value), or empty
97#[derive(Clone)]
98enum TrieNode<K, V> {
99 Internal(Box<InternalNode<K, V>>),
100 External(K, V),
101 Nothing,
102}
103
104impl<K: Chunk, V: PartialEq> PartialEq for Map<K, V> {
105 fn eq(&self, other: &Map<K, V>) -> bool {
106 self.len() == other.len() && self.iter().zip(other.iter()).all(|(a, b)| a == b)
107 }
108}
109
110impl<K: Chunk, V: Eq> Eq for Map<K, V> {}
111
112impl<K: PartialOrd + Chunk, V: PartialOrd> PartialOrd for Map<K, V> {
113 #[inline]
114 fn partial_cmp(&self, other: &Map<K, V>) -> Option<Ordering> {
115 self.iter().partial_cmp(other.iter())
116 }
117}
118
119impl<K: Ord + Chunk, V: Ord> Ord for Map<K, V> {
120 #[inline]
121 fn cmp(&self, other: &Map<K, V>) -> Ordering {
122 self.iter().cmp(other.iter())
123 }
124}
125
126impl<K: Debug + Chunk, V: Debug> Debug for Map<K, V> {
127 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
128 f.debug_map().entries(self.iter()).finish()
129 }
130}
131
132impl<K, V> Default for Map<K, V> {
133 #[inline]
134 fn default() -> Map<K, V> {
135 Map::new()
136 }
137}
138
139impl<K, V> Map<K, V> {
140 /// Creates an empty map.
141 ///
142 /// # Examples
143 ///
144 /// ```
145 /// let mut map: trie::Map<usize, &str> = trie::Map::new();
146 /// ```
147 #[inline]
148 pub fn new() -> Self {
149 Map {
150 root: InternalNode::new(),
151 length: 0,
152 }
153 }
154}
155
156impl<K: Chunk, V> Map<K, V> {
157 /// Visits all key-value pairs in reverse order. Aborts traversal when `f` returns `false`.
158 /// Returns `true` if `f` returns `true` for all elements.
159 ///
160 /// # Examples
161 ///
162 /// ```
163 /// let map: trie::Map<usize, &str> = [(1, "a"), (2, "b"), (3, "c")].iter().cloned().collect();
164 ///
165 /// let mut vec = vec![];
166 /// assert_eq!(true, map.each_reverse(|&key, &value| { vec.push((key, value)); true }));
167 /// assert_eq!(vec, [(3, "c"), (2, "b"), (1, "a")]);
168 ///
169 /// // Stop when we reach 2
170 /// let mut vec = vec![];
171 /// assert_eq!(false, map.each_reverse(|&key, &value| { vec.push(value); key != 2 }));
172 /// assert_eq!(vec, ["c", "b"]);
173 /// ```
174 #[inline]
175 pub fn each_reverse<'a, F>(&'a self, mut f: F) -> bool
176 where
177 F: FnMut(&K, &'a V) -> bool,
178 {
179 self.root.each_reverse(&mut f)
180 }
181
182 /// Gets an iterator visiting all keys in ascending order by the keys.
183 /// The iterator's element type is `usize`.
184 pub fn keys(&self) -> Keys<'_, K, V> {
185 Keys(self.iter())
186 }
187
188 /// Gets an iterator visiting all values in ascending order by the keys.
189 /// The iterator's element type is `&'r T`.
190 pub fn values(&self) -> Values<'_, K, V> {
191 Values(self.iter())
192 }
193
194 /// Gets an iterator over the key-value pairs in the map, ordered by keys.
195 ///
196 /// # Examples
197 ///
198 /// ```
199 /// let map: trie::Map<usize, &str> = [(3, "c"), (1, "a"), (2, "b")].iter().cloned().collect();
200 ///
201 /// for (key, value) in map.iter() {
202 /// println!("{}: {}", key, value);
203 /// }
204 /// ```
205 pub fn iter(&self) -> Iter<'_, K, V> {
206 let mut iter = unsafe { Iter::new() };
207 iter.stack[0] = self.root.children.iter();
208 iter.length = 1;
209 iter.remaining = self.length;
210
211 iter
212 }
213
214 /// Gets an iterator over the key-value pairs in the map, with the
215 /// ability to mutate the values.
216 ///
217 /// # Examples
218 ///
219 /// ```
220 /// let mut map: trie::Map<usize, i32> = [(1, 2), (2, 4), (3, 6)].iter().cloned().collect();
221 ///
222 /// for (&key, value) in map.iter_mut() {
223 /// *value = -(key as i32);
224 /// }
225 ///
226 /// assert_eq!(map.get(&1), Some(&-1));
227 /// assert_eq!(map.get(&2), Some(&-2));
228 /// assert_eq!(map.get(&3), Some(&-3));
229 /// ```
230 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
231 let mut iter = unsafe { IterMut::new() };
232 iter.stack[0] = self.root.children.iter_mut();
233 iter.length = 1;
234 iter.remaining = self.length;
235
236 iter
237 }
238
239 /// Return the number of elements in the map.
240 ///
241 /// # Examples
242 ///
243 /// ```
244 /// let mut a = trie::Map::new();
245 /// assert_eq!(a.len(), 0);
246 /// a.insert(1, "a");
247 /// assert_eq!(a.len(), 1);
248 /// ```
249 #[inline]
250 pub fn len(&self) -> usize {
251 self.length
252 }
253
254 /// Return true if the map contains no elements.
255 ///
256 /// # Examples
257 ///
258 /// ```
259 /// let mut a = trie::Map::new();
260 /// assert!(a.is_empty());
261 /// a.insert(1, "a");
262 /// assert!(!a.is_empty());
263 /// ```
264 #[inline]
265 pub fn is_empty(&self) -> bool {
266 self.len() == 0
267 }
268
269 /// Clears the map, removing all values.
270 ///
271 /// # Examples
272 ///
273 /// ```
274 /// let mut a = trie::Map::new();
275 /// a.insert(1, "a");
276 /// a.clear();
277 /// assert!(a.is_empty());
278 /// ```
279 #[inline]
280 pub fn clear(&mut self) {
281 self.root = InternalNode::new();
282 self.length = 0;
283 }
284
285 /// Returns a reference to the value corresponding to the key.
286 ///
287 /// # Examples
288 ///
289 /// ```
290 /// let mut map = trie::Map::new();
291 /// map.insert(1, "a");
292 /// assert_eq!(map.get(&1), Some(&"a"));
293 /// assert_eq!(map.get(&2), None);
294 /// ```
295 #[inline]
296 pub fn get<Q>(&self, key: &Q) -> Option<&V>
297 where
298 K: Borrow<Q>,
299 Q: Chunk,
300 {
301 let mut node = &self.root;
302 let mut idx = 0;
303 loop {
304 match node.children[key.chunk(idx)] {
305 Internal(ref x) => node = &**x,
306 External(ref stored, ref value) => {
307 if stored.borrow() == key {
308 return Some(value);
309 } else {
310 return None;
311 }
312 }
313 Nothing => return None,
314 }
315 idx += 1;
316 }
317 }
318
319 /// Returns true if the map contains a value for the specified key.
320 ///
321 /// # Examples
322 ///
323 /// ```
324 /// let mut map = trie::Map::new();
325 /// map.insert(1, "a");
326 /// assert_eq!(map.contains_key(&1), true);
327 /// assert_eq!(map.contains_key(&2), false);
328 /// ```
329 #[inline]
330 pub fn contains_key<Q>(&self, key: &Q) -> bool
331 where
332 K: Borrow<Q>,
333 Q: Chunk,
334 {
335 self.get(key).is_some()
336 }
337
338 /// Returns a mutable reference to the value corresponding to the key.
339 ///
340 /// # Examples
341 ///
342 /// ```
343 /// let mut map = trie::Map::new();
344 /// map.insert(1, "a");
345 /// match map.get_mut(&1) {
346 /// Some(x) => *x = "b",
347 /// None => (),
348 /// }
349 /// assert_eq!(map[&1], "b");
350 /// ```
351 #[inline]
352 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
353 where
354 K: Borrow<Q>,
355 Q: Chunk,
356 {
357 find_mut(&mut self.root.children[key.chunk(0)], key, 1)
358 }
359
360 /// Inserts a key-value pair from the map. If the key already had a value
361 /// present in the map, that value is returned. Otherwise, `None` is returned.
362 ///
363 /// # Examples
364 ///
365 /// ```
366 /// let mut map = trie::Map::new();
367 /// assert_eq!(map.insert(37, "a"), None);
368 /// assert_eq!(map.is_empty(), false);
369 ///
370 /// map.insert(37, "b");
371 /// assert_eq!(map.insert(37, "c"), Some("b"));
372 /// assert_eq!(map[&37], "c");
373 /// ```
374 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
375 let (_, old_val) = insert(
376 &mut self.root.count,
377 &mut self.root.children[key.chunk(0)],
378 key,
379 value,
380 1,
381 );
382 if old_val.is_none() {
383 self.length += 1
384 }
385 old_val
386 }
387
388 /// Removes a key from the map, returning the value at the key if the key
389 /// was previously in the map.
390 ///
391 /// # Examples
392 ///
393 /// ```
394 /// let mut map = trie::Map::new();
395 /// map.insert(1, "a");
396 /// assert_eq!(map.remove(&1), Some("a"));
397 /// assert_eq!(map.remove(&1), None);
398 /// ```
399 pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
400 where
401 K: Borrow<Q>,
402 Q: Chunk,
403 {
404 let ret = remove(
405 &mut self.root.count,
406 &mut self.root.children[key.chunk(0)],
407 key,
408 1,
409 );
410 if ret.is_some() {
411 self.length -= 1
412 }
413 ret
414 }
415}
416
417macro_rules! bound {
418 (
419 $iterator_name:ident,
420 // the current treemap
421 self = $this:expr,
422 // the key to look for
423 key = $key:expr,
424 // are we looking at the upper bound?
425 is_upper = $upper:expr,
426
427 // method name for iterating.
428 iter = $iter:ident,
429
430 // this is just an optional mut, but there's no 0-or-1 repeats yet.
431 mutability = ($($mut_:tt)*),
432 const = ($($const_:tt)*)
433 ) => {
434 {
435 // We need an unsafe pointer here because we are borrowing
436 // mutable references to the internals of each of these
437 // mutable nodes, while still using the outer node.
438 //
439 // However, we're allowed to flaunt rustc like this because we
440 // never actually modify the "shape" of the nodes. The only
441 // place that mutation is can actually occur is of the actual
442 // values of the map (as the return value of the
443 // iterator), i.e. we can never cause a deallocation of any
444 // InternalNodes so the raw pointer is always valid.
445 let this = $this;
446 let mut node = & $($mut_)* this.root as *$($mut_)* $($const_)* InternalNode<K, V>;
447
448 let key = $key;
449
450 let mut it = unsafe {$iterator_name::new()};
451 // everything else is zero'd, as we want.
452 it.remaining = this.length;
453
454 // this addr is necessary for the `Internal` pattern.
455 loop {
456 let children = unsafe { & $($mut_)* (*node).children };
457 // it.length is the current depth in the iterator and the
458 // current depth through the `usize` key we've traversed.
459 let child_id = key.chunk(it.length);
460 let (slice_idx, ret) = match & $($mut_)* children[child_id] {
461 & $($mut_)* Internal(ref $($mut_)* n) => {
462 node = (& $($mut_)* **n) as *$($mut_)* $($const_)* _;
463 (child_id + 1, false)
464 }
465 & $($mut_)* External(ref stored, _) => {
466 (if stored < key || ($upper && stored == key) {
467 child_id + 1
468 } else {
469 child_id
470 }, true)
471 }
472 & $($mut_)* Nothing => {
473 (child_id + 1, true)
474 }
475 };
476 // push to the stack.
477 it.stack[it.length] = children[slice_idx..].$iter();
478 it.length += 1;
479 if ret { break }
480 }
481
482 it
483 }
484 }
485}
486
487impl<K: Chunk + PartialOrd, V> Map<K, V> {
488 // If `upper` is true then returns upper_bound else returns lower_bound.
489 #[inline]
490 fn bound(&self, key: &K, upper: bool) -> Range<'_, K, V> {
491 Range(bound!(Iter, self = self,
492 key = key, is_upper = upper,
493 iter = iter,
494 mutability = (), const = (const)))
495 }
496
497 /// Gets an iterator pointing to the first key-value pair whose key is not less than `key`.
498 /// If all keys in the map are less than `key` an empty iterator is returned.
499 ///
500 /// # Examples
501 ///
502 /// ```
503 /// let map: trie::Map<usize, &str> = [(2, "a"), (4, "b"), (6, "c")].iter().cloned().collect();
504 ///
505 /// assert_eq!(map.lower_bound(&4).next(), Some((&4, &"b")));
506 /// assert_eq!(map.lower_bound(&5).next(), Some((&6, &"c")));
507 /// assert_eq!(map.lower_bound(&10).next(), None);
508 /// ```
509 pub fn lower_bound(&self, key: &K) -> Range<'_, K, V> {
510 self.bound(key, false)
511 }
512
513 /// Gets an iterator pointing to the first key-value pair whose key is greater than `key`.
514 /// If all keys in the map are not greater than `key` an empty iterator is returned.
515 ///
516 /// # Examples
517 ///
518 /// ```
519 /// let map: trie::Map<usize, &str> = [(2, "a"), (4, "b"), (6, "c")].iter().cloned().collect();
520 ///
521 /// assert_eq!(map.upper_bound(&4).next(), Some((&6, &"c")));
522 /// assert_eq!(map.upper_bound(&5).next(), Some((&6, &"c")));
523 /// assert_eq!(map.upper_bound(&10).next(), None);
524 /// ```
525 pub fn upper_bound(&self, key: &K) -> Range<'_, K, V> {
526 self.bound(key, true)
527 }
528 // If `upper` is true then returns upper_bound else returns lower_bound.
529 #[inline]
530 fn bound_mut(&mut self, key: &K, upper: bool) -> RangeMut<'_, K, V> {
531 RangeMut(bound!(IterMut, self = self,
532 key = key, is_upper = upper,
533 iter = iter_mut,
534 mutability = (mut), const = ()))
535 }
536
537 /// Gets an iterator pointing to the first key-value pair whose key is not less than `key`.
538 /// If all keys in the map are less than `key` an empty iterator is returned.
539 ///
540 /// # Examples
541 ///
542 /// ```
543 /// let mut map: trie::Map<usize, &str> = [(2, "a"), (4, "b"), (6, "c")].iter().cloned().collect();
544 ///
545 /// assert_eq!(map.lower_bound_mut(&4).next(), Some((&4, &mut "b")));
546 /// assert_eq!(map.lower_bound_mut(&5).next(), Some((&6, &mut "c")));
547 /// assert_eq!(map.lower_bound_mut(&10).next(), None);
548 ///
549 /// for (_key, value) in map.lower_bound_mut(&4) {
550 /// *value = "changed";
551 /// }
552 ///
553 /// assert_eq!(map.get(&2), Some(&"a"));
554 /// assert_eq!(map.get(&4), Some(&"changed"));
555 /// assert_eq!(map.get(&6), Some(&"changed"));
556 /// ```
557 pub fn lower_bound_mut(&mut self, key: &K) -> RangeMut<'_, K, V> {
558 self.bound_mut(key, false)
559 }
560
561 /// Gets an iterator pointing to the first key-value pair whose key is greater than `key`.
562 /// If all keys in the map are not greater than `key` an empty iterator is returned.
563 ///
564 /// # Examples
565 ///
566 /// ```
567 /// let mut map: trie::Map<usize, &str> = [(2, "a"), (4, "b"), (6, "c")].iter().cloned().collect();
568 ///
569 /// assert_eq!(map.upper_bound_mut(&4).next(), Some((&6, &mut "c")));
570 /// assert_eq!(map.upper_bound_mut(&5).next(), Some((&6, &mut "c")));
571 /// assert_eq!(map.upper_bound_mut(&10).next(), None);
572 ///
573 /// for (_key, value) in map.upper_bound_mut(&4) {
574 /// *value = "changed";
575 /// }
576 ///
577 /// assert_eq!(map.get(&2), Some(&"a"));
578 /// assert_eq!(map.get(&4), Some(&"b"));
579 /// assert_eq!(map.get(&6), Some(&"changed"));
580 /// ```
581 pub fn upper_bound_mut(&mut self, key: &K) -> RangeMut<'_, K, V> {
582 self.bound_mut(key, true)
583 }
584}
585
586impl<K: Chunk, V> iter::FromIterator<(K, V)> for Map<K, V> {
587 fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Map<K, V> {
588 let mut map = Map::new();
589 map.extend(iter);
590 map
591 }
592}
593
594impl<K: Chunk, V> Extend<(K, V)> for Map<K, V> {
595 fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
596 for (k, v) in iter {
597 self.insert(k, v);
598 }
599 }
600}
601
602impl<K: Chunk + Hash, V: Hash> Hash for Map<K, V> {
603 fn hash<H: Hasher>(&self, state: &mut H) {
604 for elt in self.iter() {
605 elt.hash(state);
606 }
607 }
608}
609
610impl<'a, Q, K, V> ops::Index<&'a Q> for Map<K, V>
611where
612 K: Borrow<Q> + Chunk,
613 Q: Chunk + 'a,
614{
615 type Output = V;
616 #[inline]
617 fn index(&self, i: &'a Q) -> &V {
618 self.get(i).expect("key not present")
619 }
620}
621
622impl<'a, Q, K, V> ops::IndexMut<&'a Q> for Map<K, V>
623where
624 K: Borrow<Q> + Chunk,
625 Q: Chunk + 'a,
626{
627 #[inline]
628 fn index_mut(&mut self, i: &'a Q) -> &mut V {
629 self.get_mut(i).expect("key not present")
630 }
631}
632
633impl<K, V> InternalNode<K, V> {
634 #[inline]
635 fn new() -> Self {
636 InternalNode {
637 count: 0,
638 children: [const { Nothing }; SIZE],
639 }
640 }
641}
642
643impl<K, V> InternalNode<K, V> {
644 fn each_reverse<'a, F>(&'a self, f: &mut F) -> bool
645 where
646 F: FnMut(&'a K, &'a V) -> bool,
647 {
648 for elt in self.children.iter().rev() {
649 match *elt {
650 Internal(ref x) => {
651 if !x.each_reverse(f) {
652 return false;
653 }
654 }
655 External(ref k, ref v) => {
656 if !f(k, v) {
657 return false;
658 }
659 }
660 Nothing => (),
661 }
662 }
663 true
664 }
665}
666
667fn find_mut<'a, K, Q: Chunk, V>(
668 child: &'a mut TrieNode<K, V>,
669 key: &Q,
670 idx: usize,
671) -> Option<&'a mut V>
672where
673 K: Borrow<Q> + Chunk,
674{
675 match *child {
676 External(ref stored, ref mut value) if stored.borrow() == key => Some(value),
677 External(..) => None,
678 Internal(ref mut x) => find_mut(&mut x.children[key.chunk(idx)], key, idx + 1),
679 Nothing => None,
680 }
681}
682
683/// Inserts a new node for the given key and value, at or below `start_node`.
684///
685/// The index (`idx`) is the index of the next node, such that the start node
686/// was accessed via parent.children[chunk(key, idx - 1)].
687///
688/// The count is the external node counter for the start node's parent,
689/// which will be incremented only if `start_node` is transformed into a *new* external node.
690///
691/// Returns a mutable reference to the inserted value and an optional previous value.
692fn insert<'a, K: Chunk, V>(
693 count: &mut usize,
694 start_node: &'a mut TrieNode<K, V>,
695 key: K,
696 value: V,
697 idx: usize,
698) -> (&'a mut V, Option<V>) {
699 // We branch twice to avoid having to do the `replace` when we
700 // don't need to; this is much faster, especially for keys that
701 // have long shared prefixes.
702
703 let mut hack = false;
704 match *start_node {
705 Nothing => {
706 *count += 1;
707 *start_node = External(key, value);
708 match *start_node {
709 External(_, ref mut value_ref) => return (value_ref, None),
710 _ => unreachable!(),
711 }
712 }
713 Internal(ref mut x) => {
714 let x = &mut **x;
715 return insert(
716 &mut x.count,
717 &mut x.children[key.chunk(idx)],
718 key,
719 value,
720 idx + 1,
721 );
722 }
723 External(ref stored_key, _) if stored_key == &key => {
724 hack = true;
725 }
726 _ => {}
727 }
728
729 if !hack {
730 // Conflict, an external node with differing keys.
731 // We replace the old node by an internal one, then re-insert the two values beneath it.
732 match mem::replace(start_node, Internal(Box::new(InternalNode::new()))) {
733 External(stored_key, stored_value) => {
734 match *start_node {
735 Internal(ref mut new_node) => {
736 let new_node = &mut **new_node;
737 // Re-insert the old value.
738 // TODO: remove recursive call.
739 insert(
740 &mut new_node.count,
741 &mut new_node.children[stored_key.chunk(idx)],
742 stored_key,
743 stored_value,
744 idx + 1,
745 );
746
747 // Insert the new value, and return a reference to it directly.
748 return insert(
749 &mut new_node.count,
750 &mut new_node.children[key.chunk(idx)],
751 key,
752 value,
753 idx + 1,
754 );
755 }
756 // Value that was just copied disappeared.
757 _ => unreachable!(),
758 }
759 }
760 // Logic error in previous match.
761 _ => unreachable!(),
762 }
763 }
764
765 if let External(_, ref mut stored_value) = *start_node {
766 // Swap in the new value and return the old.
767 let old_value = mem::replace(stored_value, value);
768 return (stored_value, Some(old_value));
769 }
770
771 unreachable!();
772}
773
774fn remove<K, Q: Chunk, V>(
775 count: &mut usize,
776 child: &mut TrieNode<K, V>,
777 key: &Q,
778 idx: usize,
779) -> Option<V>
780where
781 K: Borrow<Q> + Chunk,
782{
783 let (ret, this) = match *child {
784 External(ref stored, _) if stored.borrow() == key => match mem::replace(child, Nothing) {
785 External(_, value) => (Some(value), true),
786 _ => unreachable!(),
787 },
788 External(..) => (None, false),
789 Internal(ref mut x) => {
790 let x = &mut **x;
791 let ret = remove(&mut x.count, &mut x.children[key.chunk(idx)], key, idx + 1);
792 (ret, x.count == 0)
793 }
794 Nothing => (None, false),
795 };
796
797 if this {
798 *child = Nothing;
799 *count -= 1;
800 }
801 ret
802}
803
804/// A view into a single entry in a map, which may be vacant or occupied.
805pub enum Entry<'a, K: 'a, V: 'a> {
806 /// An occupied entry.
807 Occupied(OccupiedEntry<'a, K, V>),
808 /// A vacant entry.
809 Vacant(VacantEntry<'a, K, V>),
810}
811
812impl<'a, K: Chunk, V> Entry<'a, K, V> {
813 /// Ensures a value is in the entry by inserting the default if empty, and returns
814 /// a mutable reference to the value in the entry.
815 pub fn or_insert(self, default: V) -> &'a mut V {
816 match self {
817 Occupied(entry) => entry.into_mut(),
818 Vacant(entry) => entry.insert(default),
819 }
820 }
821
822 /// Ensures a value is in the entry by inserting the result of the default function if empty,
823 /// and returns a mutable reference to the value in the entry.
824 pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
825 match self {
826 Occupied(entry) => entry.into_mut(),
827 Vacant(entry) => entry.insert(default()),
828 }
829 }
830}
831
832/// A view into an occupied entry in a map.
833pub struct OccupiedEntry<'a, K: 'a, V: 'a> {
834 search_stack: SearchStack<'a, K, V>,
835}
836
837/// A view into a vacant entry in a map.
838pub struct VacantEntry<'a, K: 'a, V: 'a> {
839 search_stack: SearchStack<'a, K, V>,
840}
841
842/// A list of nodes encoding a path from the root of a map to a node.
843///
844/// Invariants:
845/// * The last node is either `External` or `Nothing`.
846/// * Pointers at indexes less than `length` can be safely dereferenced.
847///
848/// we can only use raw pointers, because of stacked borrows.
849/// Source:
850/// https://rust-unofficial.github.io/too-many-lists/fifth-stacked-borrows.html#managing-stacked-borrows
851struct SearchStack<'a, K: 'a, V: 'a> {
852 map: *mut Map<K, V>,
853 length: usize,
854 key: K,
855 items: [*mut TrieNode<K, V>; MAX_DEPTH],
856 phantom: PhantomData<&'a mut Map<K, V>>,
857}
858
859impl<'a, K, V> SearchStack<'a, K, V> {
860 /// Creates a new search-stack with empty entries.
861 fn new(map: *mut Map<K, V>, key: K) -> Self {
862 SearchStack {
863 map,
864 length: 0,
865 key,
866 items: [ptr::null_mut(); MAX_DEPTH],
867 phantom: PhantomData,
868 }
869 }
870
871 fn push(&mut self, node: *mut TrieNode<K, V>) {
872 self.length += 1;
873 self.items[self.length - 1] = node;
874 }
875
876 fn peek(&self) -> *mut TrieNode<K, V> {
877 self.items[self.length - 1]
878 }
879
880 fn peek_ref(&self) -> &'a mut TrieNode<K, V> {
881 let item = self.items[self.length - 1];
882 unsafe { &mut *item }
883 }
884
885 fn pop_ref(&mut self) -> &'a mut TrieNode<K, V> {
886 self.length -= 1;
887 unsafe { &mut *self.items[self.length] }
888 }
889
890 fn is_empty(&self) -> bool {
891 self.length == 0
892 }
893
894 fn get_ref(&self, idx: usize) -> &'a mut TrieNode<K, V> {
895 assert!(idx < self.length);
896 unsafe { &mut *self.items[idx] }
897 }
898}
899
900// Implementation of SearchStack creation logic.
901// Once a SearchStack has been created the Entry methods are relatively straight-forward.
902impl<K: Chunk, V> Map<K, V> {
903 /// Gets the given key's corresponding entry in the map for in-place manipulation.
904 #[inline]
905 pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
906 // Unconditionally add the corresponding node from the first layer.
907 let first_node = &mut self.root.children[key.chunk(0)] as *mut _;
908 // Create an empty search stack.
909 let mut search_stack = SearchStack::new(self, key);
910 search_stack.push(first_node);
911
912 // While no appropriate slot is found, keep descending down the Trie,
913 // adding nodes to the search stack.
914 let search_successful: bool;
915 loop {
916 match unsafe { next_child(search_stack.peek(), &search_stack.key, search_stack.length) }
917 {
918 (Some(child), _) => search_stack.push(child),
919 (None, success) => {
920 search_successful = success;
921 break;
922 }
923 }
924 }
925
926 if search_successful {
927 Occupied(OccupiedEntry { search_stack })
928 } else {
929 Vacant(VacantEntry { search_stack })
930 }
931 }
932}
933
934/// Get a mutable pointer to the next child of a node, given a key and an idx.
935///
936/// The idx is the index of the next child, such that `node` was accessed via
937/// parent.children[chunk(key, idx - 1)].
938///
939/// Returns a tuple with an optional mutable pointer to the next child, and
940/// a boolean flag to indicate whether the external key node was found.
941///
942/// This function is safe only if `node` points to a valid `TrieNode`.
943#[inline]
944unsafe fn next_child<K: Chunk, V>(
945 node: *mut TrieNode<K, V>,
946 key: &K,
947 idx: usize,
948) -> (Option<*mut TrieNode<K, V>>, bool) {
949 match unsafe { &mut *node } {
950 // If the node is internal, tell the caller to descend further.
951 &mut Internal(ref mut node_internal) => (
952 Some(&mut node_internal.children[key.chunk(idx)] as *mut _),
953 false,
954 ),
955 // If the node is external or empty, the search is complete.
956 // If the key doesn't match, node expansion will be done upon
957 // insertion. If it does match, we've found our node.
958 External(stored_key, _) if stored_key == key => (None, true),
959 External(..) | Nothing => (None, false),
960 }
961}
962
963// NB: All these methods assume a correctly constructed occupied entry (matching the given key).
964impl<'a, K, V> OccupiedEntry<'a, K, V> {
965 /// Gets a reference to the value in the entry.
966 #[inline]
967 pub fn get(&self) -> &V {
968 match *self.search_stack.peek_ref() {
969 External(_, ref value) => value,
970 // Invalid SearchStack, non-external last node.
971 _ => unreachable!(),
972 }
973 }
974
975 /// Gets a mutable reference to the value in the entry.
976 #[inline]
977 pub fn get_mut(&mut self) -> &mut V {
978 match *self.search_stack.peek_ref() {
979 External(_, ref mut value) => value,
980 // Invalid SearchStack, non-external last node.
981 _ => unreachable!(),
982 }
983 }
984
985 /// Converts the OccupiedEntry into a mutable reference to the value in the entry,
986 /// with a lifetime bound to the map itself.
987 #[inline]
988 pub fn into_mut(self) -> &'a mut V {
989 match *self.search_stack.peek_ref() {
990 External(_, ref mut value) => value,
991 // Invalid SearchStack, non-external last node.
992 _ => unreachable!(),
993 }
994 }
995
996 /// Sets the value of the entry, and returns the entry's old value.
997 #[inline]
998 pub fn insert(&mut self, value: V) -> V {
999 match *self.search_stack.peek_ref() {
1000 External(_, ref mut stored_value) => mem::replace(stored_value, value),
1001 // Invalid SearchStack, non-external last node.
1002 _ => unreachable!(),
1003 }
1004 }
1005
1006 /// Takes the value out of the entry, and returns it.
1007 #[inline]
1008 pub fn remove(self) -> V {
1009 // This function removes the external leaf-node, then unwinds the search-stack
1010 // deleting now-childless ancestors.
1011 let mut search_stack = self.search_stack;
1012
1013 // Extract the value from the leaf-node of interest.
1014 let leaf_node = mem::replace(search_stack.pop_ref(), Nothing);
1015
1016 let value = match leaf_node {
1017 External(_, value) => value,
1018 // Invalid SearchStack, non-external last node.
1019 _ => unreachable!(),
1020 };
1021
1022 // Iterate backwards through the search stack, deleting nodes if they are childless.
1023 // We compare each ancestor's parent count to 1 because each ancestor reached has just
1024 // had one of its children deleted.
1025 while !search_stack.is_empty() {
1026 let ancestor = search_stack.pop_ref();
1027 match *ancestor {
1028 Internal(ref mut internal) => {
1029 // If stopping deletion, update the child count and break.
1030 if internal.count != 1 {
1031 internal.count -= 1;
1032 break;
1033 }
1034 }
1035 // Invalid SearchStack, non-internal ancestor node.
1036 _ => unreachable!(),
1037 }
1038 *ancestor = Nothing;
1039 }
1040
1041 // Decrement the length of the entire map, for the removed node.
1042 unsafe {
1043 (*search_stack.map).length -= 1;
1044 }
1045
1046 value
1047 }
1048}
1049
1050impl<'a, K: Chunk, V> VacantEntry<'a, K, V> {
1051 /// Set the vacant entry to the given value.
1052 pub fn insert(self, value: V) -> &'a mut V {
1053 let search_stack = self.search_stack;
1054 let old_length = search_stack.length;
1055
1056 // Update the map's length for the new element.
1057 unsafe {
1058 (*search_stack.map).length += 1;
1059 }
1060
1061 // If there's only 1 node in the search stack, insert a new node below it at idx 1.
1062 if old_length == 1 {
1063 unsafe {
1064 // Note: Small hack to appease the borrow checker. Can't mutably borrow root.count
1065 let mut temp = (*search_stack.map).root.count;
1066 let (value_ref, _) = insert(
1067 &mut temp,
1068 search_stack.get_ref(0),
1069 search_stack.key,
1070 value,
1071 1,
1072 );
1073 (*search_stack.map).root.count = temp;
1074 value_ref
1075 }
1076 }
1077 // Otherwise, find the predecessor of the last stack node, and insert as normal.
1078 else {
1079 match *search_stack.get_ref(old_length - 2) {
1080 Internal(ref mut parent) => {
1081 let parent = &mut **parent;
1082 let (value_ref, _) = insert(
1083 &mut parent.count,
1084 &mut parent.children[search_stack.key.chunk(old_length - 1)],
1085 search_stack.key,
1086 value,
1087 old_length,
1088 );
1089 value_ref
1090 }
1091 // Invalid SearchStack, non-internal ancestor node.
1092 _ => unreachable!(),
1093 }
1094 }
1095 }
1096}
1097
1098/// A forward iterator over a map.
1099pub struct Iter<'a, K: 'a, V: 'a> {
1100 stack: [slice::Iter<'a, TrieNode<K, V>>; MAX_DEPTH],
1101 length: usize,
1102 remaining: usize,
1103}
1104
1105impl<'a, K, V> Clone for Iter<'a, K, V> {
1106 #[cfg(target_pointer_width = "32")]
1107 fn clone(&self) -> Iter<'a, T> {
1108 Iter {
1109 stack: [
1110 self.stack[0].clone(),
1111 self.stack[1].clone(),
1112 self.stack[2].clone(),
1113 self.stack[3].clone(),
1114 self.stack[4].clone(),
1115 self.stack[5].clone(),
1116 self.stack[6].clone(),
1117 self.stack[7].clone(),
1118 ],
1119 ..*self
1120 }
1121 }
1122
1123 #[cfg(target_pointer_width = "64")]
1124 fn clone(&self) -> Iter<'a, K, V> {
1125 Iter {
1126 stack: [
1127 self.stack[0].clone(),
1128 self.stack[1].clone(),
1129 self.stack[2].clone(),
1130 self.stack[3].clone(),
1131 self.stack[4].clone(),
1132 self.stack[5].clone(),
1133 self.stack[6].clone(),
1134 self.stack[7].clone(),
1135 self.stack[8].clone(),
1136 self.stack[9].clone(),
1137 self.stack[10].clone(),
1138 self.stack[11].clone(),
1139 self.stack[12].clone(),
1140 self.stack[13].clone(),
1141 self.stack[14].clone(),
1142 self.stack[15].clone(),
1143 ],
1144 ..*self
1145 }
1146 }
1147}
1148
1149/// A forward iterator over the key-value pairs of a map, with the
1150/// values being mutable.
1151pub struct IterMut<'a, K: 'a, V: 'a> {
1152 stack: [slice::IterMut<'a, TrieNode<K, V>>; MAX_DEPTH],
1153 length: usize,
1154 remaining: usize,
1155}
1156
1157/// A forward iterator over the keys of a map.
1158pub struct Keys<'a, K: 'a, V: 'a>(Iter<'a, K, V>);
1159
1160impl<'a, K, V> Clone for Keys<'a, K, V> {
1161 fn clone(&self) -> Keys<'a, K, V> {
1162 Keys(self.0.clone())
1163 }
1164}
1165
1166impl<'a, K, V> Iterator for Keys<'a, K, V> {
1167 type Item = &'a K;
1168 fn next(&mut self) -> Option<Self::Item> {
1169 self.0.next().map(|e| e.0)
1170 }
1171 fn size_hint(&self) -> (usize, Option<usize>) {
1172 self.0.size_hint()
1173 }
1174}
1175
1176impl<'a, K, V> ExactSizeIterator for Keys<'a, K, V> {}
1177
1178/// A forward iterator over the values of a map.
1179pub struct Values<'a, K: 'a, V: 'a>(Iter<'a, K, V>);
1180
1181impl<'a, K, V> Clone for Values<'a, K, V> {
1182 fn clone(&self) -> Values<'a, K, V> {
1183 Values(self.0.clone())
1184 }
1185}
1186
1187impl<'a, K, V> Iterator for Values<'a, K, V> {
1188 type Item = &'a V;
1189 fn next(&mut self) -> Option<Self::Item> {
1190 self.0.next().map(|e| e.1)
1191 }
1192 fn size_hint(&self) -> (usize, Option<usize>) {
1193 self.0.size_hint()
1194 }
1195}
1196
1197impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> {}
1198
1199macro_rules! iterator_impl {
1200 ($name:ident,
1201 iter = $iter:ident,
1202 mutability = ($($mut_:tt)*)) => {
1203 impl<'a, K, V> $name<'a, K, V> {
1204 // Create new zero'd iterator. We have a thin gilding of safety by
1205 // using init rather than uninit, so that the worst that can happen
1206 // from failing to initialise correctly after calling these is a
1207 // segfault.
1208 #[cfg(target_pointer_width="32")]
1209 unsafe fn new() -> Self {
1210 $name {
1211 remaining: 0,
1212 length: 0,
1213 // ick :( ... at least the compiler will tell us if we screwed up.
1214 stack: [IntoIterator::into_iter((& $($mut_)*[])),
1215 IntoIterator::into_iter((& $($mut_)*[])),
1216 IntoIterator::into_iter((& $($mut_)*[])),
1217 IntoIterator::into_iter((& $($mut_)*[])),
1218
1219 IntoIterator::into_iter((& $($mut_)*[])),
1220 IntoIterator::into_iter((& $($mut_)*[])),
1221 IntoIterator::into_iter((& $($mut_)*[])),
1222 IntoIterator::into_iter((& $($mut_)*[])),
1223 ]
1224 }
1225 }
1226
1227 #[cfg(target_pointer_width="64")]
1228 unsafe fn new() -> Self {
1229 $name {
1230 remaining: 0,
1231 length: 0,
1232 stack: [IntoIterator::into_iter(& $($mut_)*[]),
1233 IntoIterator::into_iter(& $($mut_)*[]),
1234 IntoIterator::into_iter(& $($mut_)*[]),
1235 IntoIterator::into_iter(& $($mut_)*[]),
1236
1237 IntoIterator::into_iter(& $($mut_)*[]),
1238 IntoIterator::into_iter(& $($mut_)*[]),
1239 IntoIterator::into_iter(& $($mut_)*[]),
1240 IntoIterator::into_iter(& $($mut_)*[]),
1241
1242 IntoIterator::into_iter(& $($mut_)*[]),
1243 IntoIterator::into_iter(& $($mut_)*[]),
1244 IntoIterator::into_iter(& $($mut_)*[]),
1245 IntoIterator::into_iter(& $($mut_)*[]),
1246
1247 IntoIterator::into_iter(& $($mut_)*[]),
1248 IntoIterator::into_iter(& $($mut_)*[]),
1249 IntoIterator::into_iter(& $($mut_)*[]),
1250 IntoIterator::into_iter(& $($mut_)*[]),
1251 ]
1252 }
1253 }
1254 }
1255
1256 impl<'a, K, V> Iterator for $name<'a, K, V> {
1257 type Item = (&'a K, &'a $($mut_)* V);
1258 // you might wonder why we're not even trying to act within the
1259 // rules, and are just manipulating raw pointers like there's no
1260 // such thing as invalid pointers and memory unsafety. The
1261 // reason is performance, without doing this we can get the
1262 // (now replaced) bench_iter_large microbenchmark down to about
1263 // 30000 ns/iter (using .unsafe_get to index self.stack directly, 38000
1264 // ns/iter with [] checked indexing), but this smashes that down
1265 // to 13500 ns/iter.
1266 //
1267 // Fortunately, it's still safe...
1268 //
1269 // We have an invariant that every Internal node
1270 // corresponds to one push to self.stack, and one pop,
1271 // nested appropriately. self.stack has enough storage
1272 // to store the maximum depth of Internal nodes in the
1273 // trie (8 on 32-bit platforms, 16 on 64-bit).
1274 fn next(&mut self) -> Option<Self::Item> {
1275 let start_ptr = self.stack.as_mut_ptr();
1276
1277 unsafe {
1278 // write_ptr is the next place to write to the stack.
1279 // invariant: start_ptr <= write_ptr < end of the
1280 // vector.
1281 let mut write_ptr = start_ptr.offset(self.length as isize);
1282 while write_ptr != start_ptr {
1283 // indexing back one is safe, since write_ptr >
1284 // start_ptr now.
1285 match (*write_ptr.offset(-1)).next() {
1286 // exhausted this iterator (i.e. finished this
1287 // Internal node), so pop from the stack.
1288 //
1289 // don't bother clearing the memory, because the
1290 // next time we use it we'll've written to it
1291 // first.
1292 None => write_ptr = write_ptr.offset(-1),
1293 Some(child) => {
1294 match *child {
1295 Internal(ref $($mut_)* node) => {
1296 // going down a level, so push
1297 // to the stack (this is the
1298 // write referenced above)
1299 *write_ptr = node.children.$iter();
1300 write_ptr = write_ptr.offset(1);
1301 }
1302 External(ref key, ref $($mut_)* value) => {
1303 self.remaining -= 1;
1304 // store the new length of the
1305 // stack, based on our current
1306 // position.
1307 self.length = (write_ptr as usize
1308 - start_ptr as usize) /
1309 mem::size_of::<slice::Iter<'_, TrieNode<K, V>>>();
1310
1311 return Some((key, value));
1312 }
1313 Nothing => {}
1314 }
1315 }
1316 }
1317 }
1318 }
1319 return None;
1320 }
1321
1322 #[inline]
1323 fn size_hint(&self) -> (usize, Option<usize>) {
1324 (self.remaining, Some(self.remaining))
1325 }
1326 }
1327
1328 impl<'a, K, V> ExactSizeIterator for $name<'a, K, V> {
1329 fn len(&self) -> usize { self.remaining }
1330 }
1331 }
1332}
1333
1334iterator_impl! { Iter, iter = iter, mutability = () }
1335iterator_impl! { IterMut, iter = iter_mut, mutability = (mut) }
1336
1337/// A bounded forward iterator over a map.
1338pub struct Range<'a, K: 'a, V: 'a>(Iter<'a, K, V>);
1339
1340impl<'a, K, V> Clone for Range<'a, K, V> {
1341 fn clone(&self) -> Range<'a, K, V> {
1342 Range(self.0.clone())
1343 }
1344}
1345
1346impl<'a, K, V> Iterator for Range<'a, K, V> {
1347 type Item = (&'a K, &'a V);
1348 fn next(&mut self) -> Option<Self::Item> {
1349 self.0.next()
1350 }
1351 fn size_hint(&self) -> (usize, Option<usize>) {
1352 (0, Some(self.0.remaining))
1353 }
1354}
1355
1356/// A bounded forward iterator over the key-value pairs of a map, with the
1357/// values being mutable.
1358pub struct RangeMut<'a, K: 'a, V: 'a>(IterMut<'a, K, V>);
1359
1360impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
1361 type Item = (&'a K, &'a mut V);
1362 fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
1363 self.0.next()
1364 }
1365 fn size_hint(&self) -> (usize, Option<usize>) {
1366 (0, Some(self.0.remaining))
1367 }
1368}
1369
1370impl<'a, K: Chunk, V> IntoIterator for &'a Map<K, V> {
1371 type Item = (&'a K, &'a V);
1372 type IntoIter = Iter<'a, K, V>;
1373 fn into_iter(self) -> Iter<'a, K, V> {
1374 self.iter()
1375 }
1376}
1377
1378impl<'a, K: Chunk, V> IntoIterator for &'a mut Map<K, V> {
1379 type Item = (&'a K, &'a mut V);
1380 type IntoIter = IterMut<'a, K, V>;
1381 fn into_iter(self) -> IterMut<'a, K, V> {
1382 self.iter_mut()
1383 }
1384}
1385
1386#[cfg(test)]
1387mod test {
1388 use std::collections::hash_map::DefaultHasher;
1389 use std::hash::{Hash, Hasher};
1390 use std::hint::black_box;
1391
1392 use super::Entry::*;
1393 use super::TrieNode::*;
1394 use super::{InternalNode, Map};
1395
1396 fn check_integrity<K, V>(trie: &InternalNode<K, V>) {
1397 assert!(trie.count != 0);
1398
1399 let mut sum = 0;
1400
1401 for x in trie.children.iter() {
1402 match *x {
1403 Nothing => (),
1404 Internal(ref y) => {
1405 check_integrity(&**y);
1406 sum += 1
1407 }
1408 External(_, _) => sum += 1,
1409 }
1410 }
1411
1412 assert_eq!(sum, trie.count);
1413 }
1414
1415 #[test]
1416 fn test_find_mut() {
1417 let mut m = Map::new();
1418 assert!(m.insert(1, 12).is_none());
1419 assert!(m.insert(2, 8).is_none());
1420 assert!(m.insert(5, 14).is_none());
1421 let new = 100;
1422 match m.get_mut(&5) {
1423 None => panic!(),
1424 Some(x) => *x = new,
1425 }
1426 assert_eq!(m.get(&5), Some(&new));
1427 }
1428
1429 #[test]
1430 fn test_find_mut_missing() {
1431 let mut m = Map::new();
1432 assert!(m.get_mut(&0).is_none());
1433 assert!(m.insert(1, 12).is_none());
1434 assert!(m.get_mut(&0).is_none());
1435 assert!(m.insert(2, 8).is_none());
1436 assert!(m.get_mut(&0).is_none());
1437 }
1438
1439 #[test]
1440 fn test_step() {
1441 let mut trie = Map::new();
1442 let n = 300;
1443
1444 for x in (1..n).step_by(2) {
1445 assert!(trie.insert(x, x + 1).is_none());
1446 assert!(trie.contains_key(&x));
1447 check_integrity(&trie.root);
1448 }
1449
1450 for x in (0..n).step_by(2) {
1451 assert!(!trie.contains_key(&x));
1452 assert!(trie.insert(x, x + 1).is_none());
1453 check_integrity(&trie.root);
1454 }
1455
1456 for x in 0..n {
1457 assert!(trie.contains_key(&x));
1458 assert!(trie.insert(x, x + 1).is_some());
1459 check_integrity(&trie.root);
1460 }
1461
1462 for x in (1..n).step_by(2) {
1463 assert!(trie.remove(&x).is_some());
1464 assert!(!trie.contains_key(&x));
1465 check_integrity(&trie.root);
1466 }
1467
1468 for x in (0..n).step_by(2) {
1469 assert!(trie.contains_key(&x));
1470 assert!(trie.insert(x, x + 1).is_some());
1471 check_integrity(&trie.root);
1472 }
1473 }
1474
1475 #[test]
1476 fn test_each_reverse() {
1477 let mut m = Map::new();
1478
1479 assert!(m.insert(3, 6).is_none());
1480 assert!(m.insert(0, 0).is_none());
1481 assert!(m.insert(4, 8).is_none());
1482 assert!(m.insert(2, 4).is_none());
1483 assert!(m.insert(1, 2).is_none());
1484
1485 let mut n = 5;
1486 let mut vec: Vec<&i32> = vec![];
1487 m.each_reverse(|k, v| {
1488 n -= 1;
1489 assert_eq!(*k, n);
1490 vec.push(v);
1491 true
1492 });
1493 assert_eq!(vec, [&8, &6, &4, &2, &0]);
1494 }
1495
1496 #[test]
1497 fn test_each_reverse_break() {
1498 let mut m = Map::new();
1499
1500 for x in (usize::MAX - 10000..usize::MAX).rev() {
1501 m.insert(x, x / 2);
1502 }
1503
1504 let mut n = usize::MAX - 1;
1505 m.each_reverse(|k, v| {
1506 if n == usize::MAX - 5000 {
1507 false
1508 } else {
1509 assert!(n > usize::MAX - 5000);
1510
1511 assert_eq!(*k, n);
1512 assert_eq!(*v, n / 2);
1513 n -= 1;
1514 true
1515 }
1516 });
1517 }
1518
1519 #[test]
1520 fn test_insert() {
1521 let mut m = Map::new();
1522 assert_eq!(m.insert(1, 2), None);
1523 assert_eq!(m.insert(1, 3), Some(2));
1524 assert_eq!(m.insert(1, 4), Some(3));
1525 }
1526
1527 #[test]
1528 fn test_remove() {
1529 let mut m = Map::new();
1530 m.insert(1, 2);
1531 assert_eq!(m.remove(&1), Some(2));
1532 assert_eq!(m.remove(&1), None);
1533 }
1534
1535 #[test]
1536 fn test_from_iter() {
1537 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1538
1539 let map: Map<usize, i32> = xs.iter().cloned().collect();
1540
1541 for &(k, v) in xs.iter() {
1542 assert_eq!(map.get(&k), Some(&v));
1543 }
1544 }
1545
1546 #[test]
1547 fn test_keys() {
1548 let vec = [(1, 'a'), (2, 'b'), (3, 'c')];
1549 let map: Map<usize, _> = vec.iter().cloned().collect();
1550 let keys: Vec<_> = map.keys().collect();
1551 assert_eq!(keys.len(), 3);
1552 assert!(keys.contains(&&1));
1553 assert!(keys.contains(&&2));
1554 assert!(keys.contains(&&3));
1555 }
1556
1557 #[test]
1558 fn test_values() {
1559 let vec = [(1, 'a'), (2, 'b'), (3, 'c')];
1560 let map: Map<usize, _> = vec.iter().cloned().collect();
1561 let values: Vec<_> = map.values().cloned().collect();
1562 assert_eq!(values.len(), 3);
1563 assert!(values.contains(&'a'));
1564 assert!(values.contains(&'b'));
1565 assert!(values.contains(&'c'));
1566 }
1567
1568 #[test]
1569 fn test_iteration() {
1570 let empty_map: Map<usize, usize> = Map::new();
1571 assert_eq!(empty_map.iter().next(), None);
1572
1573 let first = usize::MAX - 10000;
1574 let last = usize::MAX;
1575
1576 let mut map = Map::new();
1577 for x in (first..last).rev() {
1578 map.insert(x, x / 2);
1579 }
1580
1581 let mut i = 0;
1582 for (&k, &v) in map.iter() {
1583 assert_eq!(k, first + i);
1584 assert_eq!(v, k / 2);
1585 i += 1;
1586 }
1587 assert_eq!(i, last - first);
1588 }
1589
1590 #[test]
1591 fn test_mut_iter() {
1592 let mut empty_map: Map<usize, usize> = Map::new();
1593 assert!(empty_map.iter_mut().next().is_none());
1594
1595 let first = usize::MAX - 10000;
1596 let last = usize::MAX;
1597
1598 let mut map = Map::new();
1599 for x in (first..last).rev() {
1600 map.insert(x, x / 2);
1601 }
1602
1603 let mut i = 0;
1604 for (&k, v) in map.iter_mut() {
1605 assert_eq!(k, first + i);
1606 *v -= k / 2;
1607 i += 1;
1608 }
1609 assert_eq!(i, last - first);
1610
1611 assert!(map.iter().all(|(_, &v)| v == 0));
1612 }
1613
1614 #[test]
1615 fn test_bound() {
1616 let empty_map: Map<usize, usize> = Map::new();
1617 assert_eq!(empty_map.lower_bound(&0).next(), None);
1618 assert_eq!(empty_map.upper_bound(&0).next(), None);
1619
1620 let last = 999;
1621 let step = 3;
1622 let value = 42;
1623
1624 let mut map: Map<usize, usize> = Map::new();
1625 for x in (0..last).step_by(step) {
1626 assert!(x % step == 0);
1627 map.insert(x, value);
1628 }
1629
1630 for i in 0..last - step {
1631 let mut lb = map.lower_bound(&i);
1632 let mut ub = map.upper_bound(&i);
1633 let next_key = i - i % step + step;
1634 let next_pair = (&next_key, &value);
1635 if i % step == 0 {
1636 assert_eq!(lb.next(), Some((&i, &value)));
1637 } else {
1638 assert_eq!(lb.next(), Some(next_pair));
1639 }
1640 assert_eq!(ub.next(), Some(next_pair));
1641 }
1642
1643 let mut lb = map.lower_bound(&(last - step));
1644 assert_eq!(lb.next(), Some((&(last - step), &value)));
1645 let mut ub = map.upper_bound(&(last - step));
1646 assert_eq!(ub.next(), None);
1647
1648 for i in last - step + 1..last {
1649 let mut lb = map.lower_bound(&i);
1650 assert_eq!(lb.next(), None);
1651 let mut ub = map.upper_bound(&i);
1652 assert_eq!(ub.next(), None);
1653 }
1654 }
1655
1656 #[test]
1657 fn test_mut_bound() {
1658 let empty_map: Map<usize, usize> = Map::new();
1659 assert_eq!(empty_map.lower_bound(&0).next(), None);
1660 assert_eq!(empty_map.upper_bound(&0).next(), None);
1661
1662 let mut m_lower = Map::new();
1663 let mut m_upper = Map::new();
1664 for i in 0..100 {
1665 m_lower.insert(2 * i, 4 * i);
1666 m_upper.insert(2 * i, 4 * i);
1667 }
1668
1669 for i in 0..199 {
1670 let mut lb_it = m_lower.lower_bound_mut(&i);
1671 let (&k, v) = lb_it.next().unwrap();
1672 let lb = i + i % 2;
1673 assert_eq!(lb, k);
1674 *v -= k;
1675 }
1676
1677 for i in 0..198 {
1678 let mut ub_it = m_upper.upper_bound_mut(&i);
1679 let (&k, v) = ub_it.next().unwrap();
1680 let ub = i + 2 - i % 2;
1681 assert_eq!(ub, k);
1682 *v -= k;
1683 }
1684
1685 assert!(m_lower.lower_bound_mut(&199).next().is_none());
1686 assert!(m_upper.upper_bound_mut(&198).next().is_none());
1687
1688 assert!(m_lower.iter().all(|(_, &x)| x == 0));
1689 assert!(m_upper.iter().all(|(_, &x)| x == 0));
1690 }
1691
1692 #[test]
1693 fn test_clone() {
1694 let mut a = Map::new();
1695
1696 a.insert(1, 'a');
1697 a.insert(2, 'b');
1698 a.insert(3, 'c');
1699
1700 assert!(a.clone() == a);
1701 }
1702
1703 #[test]
1704 fn test_eq() {
1705 let mut a = Map::new();
1706 let mut b = Map::new();
1707
1708 assert!(a == b);
1709 assert!(a.insert(0, 5).is_none());
1710 assert!(a != b);
1711 assert!(b.insert(0, 4).is_none());
1712 assert!(a != b);
1713 assert!(a.insert(5, 19).is_none());
1714 assert!(a != b);
1715 assert!(b.insert(0, 5).is_some());
1716 assert!(a != b);
1717 assert!(b.insert(5, 19).is_none());
1718 assert!(a == b);
1719 }
1720
1721 #[test]
1722 fn test_lt() {
1723 let mut a = Map::new();
1724 let mut b = Map::new();
1725
1726 assert!((a >= b) && (b >= a));
1727 assert!(b.insert(2, 5).is_none());
1728 assert!(a < b);
1729 assert!(a.insert(2, 7).is_none());
1730 assert!((a >= b) && b < a);
1731 assert!(b.insert(1, 0).is_none());
1732 assert!(b < a);
1733 assert!(a.insert(0, 6).is_none());
1734 assert!(a < b);
1735 assert!(a.insert(6, 2).is_none());
1736 assert!(a < b && (b >= a));
1737 }
1738
1739 #[test]
1740 fn test_ord() {
1741 let mut a = Map::new();
1742 let mut b = Map::new();
1743
1744 assert!(a == b);
1745 assert!(a.insert(1, 1).is_none());
1746 assert!(a > b && a >= b);
1747 assert!(b < a && b <= a);
1748 assert!(b.insert(2, 2).is_none());
1749 assert!(b > a && b >= a);
1750 assert!(a < b && a <= b);
1751 }
1752
1753 #[test]
1754 fn test_hash() {
1755 fn hash<T: Hash>(t: &T) -> u64 {
1756 let mut s = DefaultHasher::new();
1757 t.hash(&mut s);
1758 s.finish()
1759 }
1760
1761 let mut x = Map::new();
1762 let mut y = Map::new();
1763
1764 assert!(hash(&x) == hash(&y));
1765 x.insert(1, 'a');
1766 x.insert(2, 'b');
1767 x.insert(3, 'c');
1768
1769 y.insert(3, 'c');
1770 y.insert(2, 'b');
1771 y.insert(1, 'a');
1772
1773 assert!(hash(&x) == hash(&y));
1774 }
1775
1776 #[test]
1777 fn test_debug() {
1778 let mut map = Map::new();
1779 let empty: Map<usize, char> = Map::new();
1780
1781 map.insert(1, 'a');
1782 map.insert(2, 'b');
1783
1784 assert_eq!(format!("{:?}", map), "{1: 'a', 2: 'b'}");
1785 assert_eq!(format!("{:?}", empty), "{}");
1786 }
1787
1788 #[test]
1789 fn test_index() {
1790 let mut map = Map::new();
1791
1792 map.insert(1, 2);
1793 map.insert(2, 1);
1794 map.insert(3, 4);
1795
1796 assert_eq!(map[&2], 1);
1797 }
1798
1799 #[test]
1800 #[should_panic]
1801 fn test_index_nonexistent() {
1802 let mut map = Map::new();
1803
1804 map.insert(1, 2);
1805 map.insert(2, 1);
1806 map.insert(3, 4);
1807
1808 black_box(map[&4]);
1809 }
1810
1811 // Number of items to insert into the map during entry tests.
1812 // The tests rely on it being even.
1813 const SQUARES_UPPER_LIM: usize = 128;
1814
1815 /// Make a map storing i^2 for i in [0, 128)
1816 fn squares_map() -> Map<usize, usize> {
1817 let mut map = Map::new();
1818 for i in 0..SQUARES_UPPER_LIM {
1819 map.insert(i, i * i);
1820 }
1821 map
1822 }
1823
1824 #[test]
1825 fn test_entry_get() {
1826 let mut map = squares_map();
1827
1828 for i in 0..SQUARES_UPPER_LIM {
1829 match map.entry(i) {
1830 Occupied(slot) => assert_eq!(slot.get(), &(i * i)),
1831 Vacant(_) => panic!("Key not found."),
1832 }
1833 }
1834 check_integrity(&map.root);
1835 }
1836
1837 #[test]
1838 fn test_entry_get_mut() {
1839 let mut map = squares_map();
1840
1841 // Change the entries to cubes.
1842 for i in 0..SQUARES_UPPER_LIM {
1843 match map.entry(i) {
1844 Occupied(mut e) => {
1845 *e.get_mut() = i * i * i;
1846 }
1847 Vacant(_) => panic!("Key not found."),
1848 }
1849 assert_eq!(map.get(&i).unwrap(), &(i * i * i));
1850 }
1851
1852 check_integrity(&map.root);
1853 }
1854
1855 #[test]
1856 fn test_entry_into_mut() {
1857 let mut map = Map::new();
1858 map.insert(3, 6);
1859
1860 let value_ref = match map.entry(3) {
1861 Occupied(e) => e.into_mut(),
1862 Vacant(_) => panic!("Entry not found."),
1863 };
1864
1865 assert_eq!(*value_ref, 6);
1866 }
1867
1868 #[test]
1869 fn test_entry_take() {
1870 let mut map = squares_map();
1871 assert_eq!(map.len(), SQUARES_UPPER_LIM);
1872
1873 // Remove every odd key, checking that the correct value is returned.
1874 for i in (1..SQUARES_UPPER_LIM).step_by(2) {
1875 match map.entry(i) {
1876 Occupied(e) => assert_eq!(e.remove(), i * i),
1877 Vacant(_) => panic!("Key not found."),
1878 }
1879 }
1880
1881 check_integrity(&map.root);
1882
1883 // Check that the values for even keys remain unmodified.
1884 for i in (0..SQUARES_UPPER_LIM).step_by(2) {
1885 assert_eq!(map.get(&i).unwrap(), &(i * i));
1886 }
1887
1888 assert_eq!(map.len(), SQUARES_UPPER_LIM / 2);
1889 }
1890
1891 #[test]
1892 fn test_occupied_entry_set() {
1893 let mut map = squares_map();
1894
1895 // Change all the entries to cubes.
1896 for i in 0..SQUARES_UPPER_LIM {
1897 match map.entry(i) {
1898 Occupied(mut e) => assert_eq!(e.insert(i * i * i), i * i),
1899 Vacant(_) => panic!("Key not found."),
1900 }
1901 assert_eq!(map.get(&i).unwrap(), &(i * i * i));
1902 }
1903 check_integrity(&map.root);
1904 }
1905
1906 #[test]
1907 fn test_vacant_entry_set() {
1908 let mut map = Map::new();
1909
1910 for i in 0..SQUARES_UPPER_LIM {
1911 match map.entry(i) {
1912 Vacant(e) => {
1913 // Insert i^2.
1914 let inserted_val = e.insert(i * i);
1915 assert_eq!(*inserted_val, i * i);
1916
1917 // Update it to i^3 using the returned mutable reference.
1918 *inserted_val = i * i * i;
1919 }
1920 _ => panic!("Non-existent key found."),
1921 }
1922 assert_eq!(map.get(&i).unwrap(), &(i * i * i));
1923 }
1924
1925 check_integrity(&map.root);
1926 assert_eq!(map.len(), SQUARES_UPPER_LIM);
1927 }
1928
1929 #[test]
1930 fn test_single_key() {
1931 let mut map = Map::new();
1932 map.insert(1, 2);
1933
1934 if let Occupied(e) = map.entry(1) {
1935 e.remove();
1936 }
1937 }
1938}