A trie in Rust
at generic 1938 lines 60 kB view raw
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}