compiler/rustc_data_structures/src/sso/map.rs

   1 use super::either_iter::EitherIter;
   2 use crate::fx::FxHashMap;
   3 use arrayvec::ArrayVec;
   4 use std::fmt;
   5 use std::hash::Hash;
   6 use std::iter::FromIterator;
   7 use std::ops::Index;
   8
   9 // For pointer-sized arguments arrays
  10 // are faster than set/map for up to 64
  11 // arguments.
  12 //
  13 // On the other hand such a big array
  14 // hurts cache performance, makes passing
  15 // sso structures around very expensive.
  16 //
  17 // Biggest performance benefit is gained
  18 // for reasonably small arrays that stay
  19 // small in vast majority of cases.
  20 //
  21 // '8' is chosen as a sane default, to be
  22 // reevaluated later.
  23 const SSO_ARRAY_SIZE: usize = 8;
  24
  25 /// Small-storage-optimized implementation of a map.
  26 ///
  27 /// Stores elements in a small array up to a certain length
  28 /// and switches to `HashMap` when that length is exceeded.
  29 //
  30 // FIXME: Implements subset of HashMap API.
  31 //
  32 // Missing HashMap API:
  33 //   all hasher-related
  34 //   try_reserve
  35 //   shrink_to (unstable)
  36 //   drain_filter (unstable)
  37 //   into_keys/into_values (unstable)
  38 //   all raw_entry-related
  39 //   PartialEq/Eq (requires sorting the array)
  40 //   Entry::or_insert_with_key
  41 //   Vacant/Occupied entries and related
  42 //
  43 // FIXME: In HashMap most methods accepting key reference
  44 // accept reference to generic `Q` where `K: Borrow<Q>`.
  45 //
  46 // However, using this approach in `HashMap::get` apparently
  47 // breaks inlining and noticeably reduces performance.
  48 //
  49 // Performance *should* be the same given that borrow is
  50 // a NOP in most cases, but in practice that's not the case.
  51 //
  52 // Further investigation is required.
  53 //
  54 // Affected methods:
  55 //   SsoHashMap::get
  56 //   SsoHashMap::get_mut
  57 //   SsoHashMap::get_entry
  58 //   SsoHashMap::get_key_value
  59 //   SsoHashMap::contains_key
  60 //   SsoHashMap::remove
  61 //   SsoHashMap::remove_entry
  62 //   Index::index
  63 //   SsoHashSet::take
  64 //   SsoHashSet::get
  65 //   SsoHashSet::remove
  66 //   SsoHashSet::contains
  67
  68 #[derive(Clone)]
  69 pub enum SsoHashMap<K, V> {
  70     Array(ArrayVec<(K, V), SSO_ARRAY_SIZE>),
  71     Map(FxHashMap<K, V>),
  72 }
  73
  74 impl<K, V> SsoHashMap<K, V> {
  75     /// Creates an empty `SsoHashMap`.
  76     #[inline]
  77     pub fn new() -> Self {
  78         SsoHashMap::Array(ArrayVec::new())
  79     }
  80
  81     /// Creates an empty `SsoHashMap` with the specified capacity.
  82     pub fn with_capacity(cap: usize) -> Self {
  83         if cap <= SSO_ARRAY_SIZE {
  84             Self::new()
  85         } else {
  86             SsoHashMap::Map(FxHashMap::with_capacity_and_hasher(cap, Default::default()))
  87         }
  88     }
  89
  90     /// Clears the map, removing all key-value pairs. Keeps the allocated memory
  91     /// for reuse.
  92     pub fn clear(&mut self) {
  93         match self {
  94             SsoHashMap::Array(array) => array.clear(),
  95             SsoHashMap::Map(map) => map.clear(),
  96         }
  97     }
  98
  99     /// Returns the number of elements the map can hold without reallocating.
 100     pub fn capacity(&self) -> usize {
 101         match self {
 102             SsoHashMap::Array(_) => SSO_ARRAY_SIZE,
 103             SsoHashMap::Map(map) => map.capacity(),
 104         }
 105     }
 106
 107     /// Returns the number of elements in the map.
 108     pub fn len(&self) -> usize {
 109         match self {
 110             SsoHashMap::Array(array) => array.len(),
 111             SsoHashMap::Map(map) => map.len(),
 112         }
 113     }
 114
 115     /// Returns `true` if the map contains no elements.
 116     pub fn is_empty(&self) -> bool {
 117         match self {
 118             SsoHashMap::Array(array) => array.is_empty(),
 119             SsoHashMap::Map(map) => map.is_empty(),
 120         }
 121     }
 122
 123     /// An iterator visiting all key-value pairs in arbitrary order.
 124     /// The iterator element type is `(&'a K, &'a V)`.
 125     #[inline]
 126     pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
 127         self.into_iter()
 128     }
 129
 130     /// An iterator visiting all key-value pairs in arbitrary order,
 131     /// with mutable references to the values.
 132     /// The iterator element type is `(&'a K, &'a mut V)`.
 133     #[inline]
 134     pub fn iter_mut(&mut self) -> impl Iterator<Item = (&'_ K, &'_ mut V)> {
 135         self.into_iter()
 136     }
 137
 138     /// An iterator visiting all keys in arbitrary order.
 139     /// The iterator element type is `&'a K`.
 140     pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
 141         match self {
 142             SsoHashMap::Array(array) => EitherIter::Left(array.iter().map(|(k, _v)| k)),
 143             SsoHashMap::Map(map) => EitherIter::Right(map.keys()),
 144         }
 145     }
 146
 147     /// An iterator visiting all values in arbitrary order.
 148     /// The iterator element type is `&'a V`.
 149     pub fn values(&self) -> impl Iterator<Item = &'_ V> {
 150         match self {
 151             SsoHashMap::Array(array) => EitherIter::Left(array.iter().map(|(_k, v)| v)),
 152             SsoHashMap::Map(map) => EitherIter::Right(map.values()),
 153         }
 154     }
 155
 156     /// An iterator visiting all values mutably in arbitrary order.
 157     /// The iterator element type is `&'a mut V`.
 158     pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
 159         match self {
 160             SsoHashMap::Array(array) => EitherIter::Left(array.iter_mut().map(|(_k, v)| v)),
 161             SsoHashMap::Map(map) => EitherIter::Right(map.values_mut()),
 162         }
 163     }
 164
 165     /// Clears the map, returning all key-value pairs as an iterator. Keeps the
 166     /// allocated memory for reuse.
 167     pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
 168         match self {
 169             SsoHashMap::Array(array) => EitherIter::Left(array.drain(..)),
 170             SsoHashMap::Map(map) => EitherIter::Right(map.drain()),
 171         }
 172     }
 173 }
 174
 175 impl<K: Eq + Hash, V> SsoHashMap<K, V> {
 176     /// Changes underlying storage from array to hashmap
 177     /// if array is full.
 178     fn migrate_if_full(&mut self) {
 179         if let SsoHashMap::Array(array) = self {
 180             if array.is_full() {
 181                 *self = SsoHashMap::Map(array.drain(..).collect());
 182             }
 183         }
 184     }
 185
 186     /// Reserves capacity for at least `additional` more elements to be inserted
 187     /// in the `SsoHashMap`. The collection may reserve more space to avoid
 188     /// frequent reallocations.
 189     pub fn reserve(&mut self, additional: usize) {
 190         match self {
 191             SsoHashMap::Array(array) => {
 192                 if SSO_ARRAY_SIZE < (array.len() + additional) {
 193                     let mut map: FxHashMap<K, V> = array.drain(..).collect();
 194                     map.reserve(additional);
 195                     *self = SsoHashMap::Map(map);
 196                 }
 197             }
 198             SsoHashMap::Map(map) => map.reserve(additional),
 199         }
 200     }
 201
 202     /// Shrinks the capacity of the map as much as possible. It will drop
 203     /// down as much as possible while maintaining the internal rules
 204     /// and possibly leaving some space in accordance with the resize policy.
 205     pub fn shrink_to_fit(&mut self) {
 206         if let SsoHashMap::Map(map) = self {
 207             if map.len() <= SSO_ARRAY_SIZE {
 208                 *self = SsoHashMap::Array(map.drain().collect());
 209             } else {
 210                 map.shrink_to_fit();
 211             }
 212         }
 213     }
 214
 215     /// Retains only the elements specified by the predicate.
 216     pub fn retain<F>(&mut self, mut f: F)
 217     where
 218         F: FnMut(&K, &mut V) -> bool,
 219     {
 220         match self {
 221             SsoHashMap::Array(array) => array.retain(|(k, v)| f(k, v)),
 222             SsoHashMap::Map(map) => map.retain(f),
 223         }
 224     }
 225
 226     /// Inserts a key-value pair into the map.
 227     ///
 228     /// If the map did not have this key present, [`None`] is returned.
 229     ///
 230     /// If the map did have this key present, the value is updated, and the old
 231     /// value is returned. The key is not updated, though; this matters for
 232     /// types that can be `==` without being identical. See the [module-level
 233     /// documentation] for more.
 234     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
 235         match self {
 236             SsoHashMap::Array(array) => {
 237                 for (k, v) in array.iter_mut() {
 238                     if *k == key {
 239                         let old_value = std::mem::replace(v, value);
 240                         return Some(old_value);
 241                     }
 242                 }
 243                 if let Err(error) = array.try_push((key, value)) {
 244                     let mut map: FxHashMap<K, V> = array.drain(..).collect();
 245                     let (key, value) = error.element();
 246                     map.insert(key, value);
 247                     *self = SsoHashMap::Map(map);
 248                 }
 249                 None
 250             }
 251             SsoHashMap::Map(map) => map.insert(key, value),
 252         }
 253     }
 254
 255     /// Removes a key from the map, returning the value at the key if the key
 256     /// was previously in the map.
 257     pub fn remove(&mut self, key: &K) -> Option<V> {
 258         match self {
 259             SsoHashMap::Array(array) => {
 260                 if let Some(index) = array.iter().position(|(k, _v)| k == key) {
 261                     Some(array.swap_remove(index).1)
 262                 } else {
 263                     None
 264                 }
 265             }
 266             SsoHashMap::Map(map) => map.remove(key),
 267         }
 268     }
 269
 270     /// Removes a key from the map, returning the stored key and value if the
 271     /// key was previously in the map.
 272     pub fn remove_entry(&mut self, key: &K) -> Option<(K, V)> {
 273         match self {
 274             SsoHashMap::Array(array) => {
 275                 if let Some(index) = array.iter().position(|(k, _v)| k == key) {
 276                     Some(array.swap_remove(index))
 277                 } else {
 278                     None
 279                 }
 280             }
 281             SsoHashMap::Map(map) => map.remove_entry(key),
 282         }
 283     }
 284
 285     /// Returns a reference to the value corresponding to the key.
 286     pub fn get(&self, key: &K) -> Option<&V> {
 287         match self {
 288             SsoHashMap::Array(array) => {
 289                 for (k, v) in array {
 290                     if k == key {
 291                         return Some(v);
 292                     }
 293                 }
 294                 None
 295             }
 296             SsoHashMap::Map(map) => map.get(key),
 297         }
 298     }
 299
 300     /// Returns a mutable reference to the value corresponding to the key.
 301     pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
 302         match self {
 303             SsoHashMap::Array(array) => {
 304                 for (k, v) in array {
 305                     if k == key {
 306                         return Some(v);
 307                     }
 308                 }
 309                 None
 310             }
 311             SsoHashMap::Map(map) => map.get_mut(key),
 312         }
 313     }
 314
 315     /// Returns the key-value pair corresponding to the supplied key.
 316     pub fn get_key_value(&self, key: &K) -> Option<(&K, &V)> {
 317         match self {
 318             SsoHashMap::Array(array) => {
 319                 for (k, v) in array {
 320                     if k == key {
 321                         return Some((k, v));
 322                     }
 323                 }
 324                 None
 325             }
 326             SsoHashMap::Map(map) => map.get_key_value(key),
 327         }
 328     }
 329
 330     /// Returns `true` if the map contains a value for the specified key.
 331     pub fn contains_key(&self, key: &K) -> bool {
 332         match self {
 333             SsoHashMap::Array(array) => array.iter().any(|(k, _v)| k == key),
 334             SsoHashMap::Map(map) => map.contains_key(key),
 335         }
 336     }
 337
 338     /// Gets the given key's corresponding entry in the map for in-place manipulation.
 339     #[inline]
 340     pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
 341         Entry { ssomap: self, key }
 342     }
 343 }
 344
 345 impl<K, V> Default for SsoHashMap<K, V> {
 346     #[inline]
 347     fn default() -> Self {
 348         Self::new()
 349     }
 350 }
 351
 352 impl<K: Eq + Hash, V> FromIterator<(K, V)> for SsoHashMap<K, V> {
 353     fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> SsoHashMap<K, V> {
 354         let mut map: SsoHashMap<K, V> = Default::default();
 355         map.extend(iter);
 356         map
 357     }
 358 }
 359
 360 impl<K: Eq + Hash, V> Extend<(K, V)> for SsoHashMap<K, V> {
 361     fn extend<I>(&mut self, iter: I)
 362     where
 363         I: IntoIterator<Item = (K, V)>,
 364     {
 365         for (key, value) in iter.into_iter() {
 366             self.insert(key, value);
 367         }
 368     }
 369
 370     #[inline]
 371     fn extend_one(&mut self, (k, v): (K, V)) {
 372         self.insert(k, v);
 373     }
 374
 375     fn extend_reserve(&mut self, additional: usize) {
 376         match self {
 377             SsoHashMap::Array(array) => {
 378                 if SSO_ARRAY_SIZE < (array.len() + additional) {
 379                     let mut map: FxHashMap<K, V> = array.drain(..).collect();
 380                     map.extend_reserve(additional);
 381                     *self = SsoHashMap::Map(map);
 382                 }
 383             }
 384             SsoHashMap::Map(map) => map.extend_reserve(additional),
 385         }
 386     }
 387 }
 388
 389 impl<'a, K, V> Extend<(&'a K, &'a V)> for SsoHashMap<K, V>
 390 where
 391     K: Eq + Hash + Copy,
 392     V: Copy,
 393 {
 394     fn extend<T: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: T) {
 395         self.extend(iter.into_iter().map(|(k, v)| (*k, *v)))
 396     }
 397
 398     #[inline]
 399     fn extend_one(&mut self, (&k, &v): (&'a K, &'a V)) {
 400         self.insert(k, v);
 401     }
 402
 403     #[inline]
 404     fn extend_reserve(&mut self, additional: usize) {
 405         Extend::<(K, V)>::extend_reserve(self, additional)
 406     }
 407 }
 408
 409 impl<K, V> IntoIterator for SsoHashMap<K, V> {
 410     type IntoIter = EitherIter<
 411         <ArrayVec<(K, V), 8> as IntoIterator>::IntoIter,
 412         <FxHashMap<K, V> as IntoIterator>::IntoIter,
 413     >;
 414     type Item = <Self::IntoIter as Iterator>::Item;
 415
 416     fn into_iter(self) -> Self::IntoIter {
 417         match self {
 418             SsoHashMap::Array(array) => EitherIter::Left(array.into_iter()),
 419             SsoHashMap::Map(map) => EitherIter::Right(map.into_iter()),
 420         }
 421     }
 422 }
 423
 424 /// adapts Item of array reference iterator to Item of hashmap reference iterator.
 425 #[inline(always)]
 426 fn adapt_array_ref_it<K, V>(pair: &'a (K, V)) -> (&'a K, &'a V) {
 427     let (a, b) = pair;
 428     (a, b)
 429 }
 430
 431 /// adapts Item of array mut reference iterator to Item of hashmap mut reference iterator.
 432 #[inline(always)]
 433 fn adapt_array_mut_it<K, V>(pair: &'a mut (K, V)) -> (&'a K, &'a mut V) {
 434     let (a, b) = pair;
 435     (a, b)
 436 }
 437
 438 impl<'a, K, V> IntoIterator for &'a SsoHashMap<K, V> {
 439     type IntoIter = EitherIter<
 440         std::iter::Map<
 441             <&'a ArrayVec<(K, V), 8> as IntoIterator>::IntoIter,
 442             fn(&'a (K, V)) -> (&'a K, &'a V),
 443         >,
 444         <&'a FxHashMap<K, V> as IntoIterator>::IntoIter,
 445     >;
 446     type Item = <Self::IntoIter as Iterator>::Item;
 447
 448     fn into_iter(self) -> Self::IntoIter {
 449         match self {
 450             SsoHashMap::Array(array) => EitherIter::Left(array.into_iter().map(adapt_array_ref_it)),
 451             SsoHashMap::Map(map) => EitherIter::Right(map.iter()),
 452         }
 453     }
 454 }
 455
 456 impl<'a, K, V> IntoIterator for &'a mut SsoHashMap<K, V> {
 457     type IntoIter = EitherIter<
 458         std::iter::Map<
 459             <&'a mut ArrayVec<(K, V), 8> as IntoIterator>::IntoIter,
 460             fn(&'a mut (K, V)) -> (&'a K, &'a mut V),
 461         >,
 462         <&'a mut FxHashMap<K, V> as IntoIterator>::IntoIter,
 463     >;
 464     type Item = <Self::IntoIter as Iterator>::Item;
 465
 466     fn into_iter(self) -> Self::IntoIter {
 467         match self {
 468             SsoHashMap::Array(array) => EitherIter::Left(array.into_iter().map(adapt_array_mut_it)),
 469             SsoHashMap::Map(map) => EitherIter::Right(map.iter_mut()),
 470         }
 471     }
 472 }
 473
 474 impl<K, V> fmt::Debug for SsoHashMap<K, V>
 475 where
 476     K: fmt::Debug,
 477     V: fmt::Debug,
 478 {
 479     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 480         f.debug_map().entries(self.iter()).finish()
 481     }
 482 }
 483
 484 impl<'a, K, V> Index<&'a K> for SsoHashMap<K, V>
 485 where
 486     K: Eq + Hash,
 487 {
 488     type Output = V;
 489
 490     #[inline]
 491     fn index(&self, key: &K) -> &V {
 492         self.get(key).expect("no entry found for key")
 493     }
 494 }
 495
 496 /// A view into a single entry in a map.
 497 pub struct Entry<'a, K, V> {
 498     ssomap: &'a mut SsoHashMap<K, V>,
 499     key: K,
 500 }
 501
 502 impl<'a, K: Eq + Hash, V> Entry<'a, K, V> {
 503     /// Provides in-place mutable access to an occupied entry before any
 504     /// potential inserts into the map.
 505     pub fn and_modify<F>(self, f: F) -> Self
 506     where
 507         F: FnOnce(&mut V),
 508     {
 509         if let Some(value) = self.ssomap.get_mut(&self.key) {
 510             f(value);
 511         }
 512         self
 513     }
 514
 515     /// Ensures a value is in the entry by inserting the default if empty, and returns
 516     /// a mutable reference to the value in the entry.
 517     #[inline]
 518     pub fn or_insert(self, value: V) -> &'a mut V {
 519         self.or_insert_with(|| value)
 520     }
 521
 522     /// Ensures a value is in the entry by inserting the result of the default function if empty,
 523     /// and returns a mutable reference to the value in the entry.
 524     pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
 525         self.ssomap.migrate_if_full();
 526         match self.ssomap {
 527             SsoHashMap::Array(array) => {
 528                 let key_ref = &self.key;
 529                 let found_index = array.iter().position(|(k, _v)| k == key_ref);
 530                 let index = if let Some(index) = found_index {
 531                     index
 532                 } else {
 533                     let index = array.len();
 534                     array.try_push((self.key, default())).unwrap();
 535                     index
 536                 };
 537                 &mut array[index].1
 538             }
 539             SsoHashMap::Map(map) => map.entry(self.key).or_insert_with(default),
 540         }
 541     }
 542
 543     /// Returns a reference to this entry's key.
 544     #[inline]
 545     pub fn key(&self) -> &K {
 546         &self.key
 547     }
 548 }
 549
 550 impl<'a, K: Eq + Hash, V: Default> Entry<'a, K, V> {
 551     /// Ensures a value is in the entry by inserting the default value if empty,
 552     /// and returns a mutable reference to the value in the entry.
 553     #[inline]
 554     pub fn or_default(self) -> &'a mut V {
 555         self.or_insert_with(Default::default)
 556     }
 557 }