SsoHashSet/SsoHashMap API greatly expanded

Now both provide almost complete API of their non-SSO counterparts.

1 files changed, 488 insertions, 20 deletions

diff --git a/compiler/rustc_data_structures/src/sso/map.rs b/compiler/rustc_data_structures/src/sso/map.rs
index c253e9d6616..258368c8ef3 100644
--- a/compiler/rustc_data_structures/src/sso/map.rs
+++ b/compiler/rustc_data_structures/src/sso/map.rs
@@ -1,32 +1,202 @@
+use super::EitherIter;
 use crate::fx::FxHashMap;
 use arrayvec::ArrayVec;
-
+use std::borrow::Borrow;
+use std::fmt;
 use std::hash::Hash;
+use std::iter::FromIterator;
+use std::ops::Index;
 
-/// Small-storage-optimized implementation of a map
-/// made specifically for caching results.
+/// Small-storage-optimized implementation of a map.
 ///
 /// Stores elements in a small array up to a certain length
 /// and switches to `HashMap` when that length is exceeded.
+///
+/// Implements subset of HashMap API.
+///
+/// Missing HashMap API:
+///   all hasher-related
+///   try_reserve (unstable)
+///   shrink_to (unstable)
+///   drain_filter (unstable)
+///   into_keys/into_values (unstable)
+///   all raw_entry-related
+///   PartialEq/Eq (requires sorting the array)
+///   Entry::or_insert_with_key (unstable)
+///   Vacant/Occupied entries and related
+#[derive(Clone)]
 pub enum SsoHashMap<K, V> {
     Array(ArrayVec<[(K, V); 8]>),
     Map(FxHashMap<K, V>),
 }
 
-impl<K: Eq + Hash, V> SsoHashMap<K, V> {
+impl<K, V> SsoHashMap<K, V> {
     /// Creates an empty `SsoHashMap`.
     pub fn new() -> Self {
         SsoHashMap::Array(ArrayVec::new())
     }
 
-    /// Inserts or updates value in the map.
-    pub fn insert(&mut self, key: K, value: V) {
+    /// Creates an empty `SsoHashMap` with the specified capacity.
+    pub fn with_capacity(cap: usize) -> Self {
+        let array = ArrayVec::new();
+        if array.capacity() >= cap {
+            SsoHashMap::Array(array)
+        } else {
+            SsoHashMap::Map(FxHashMap::with_capacity_and_hasher(cap, Default::default()))
+        }
+    }
+
+    /// Clears the map, removing all key-value pairs. Keeps the allocated memory
+    /// for reuse.
+    pub fn clear(&mut self) {
+        match self {
+            SsoHashMap::Array(array) => array.clear(),
+            SsoHashMap::Map(map) => map.clear(),
+        }
+    }
+
+    /// Returns the number of elements the map can hold without reallocating.
+    pub fn capacity(&self) -> usize {
+        match self {
+            SsoHashMap::Array(array) => array.capacity(),
+            SsoHashMap::Map(map) => map.capacity(),
+        }
+    }
+
+    /// Returns the number of elements in the map.
+    pub fn len(&self) -> usize {
+        match self {
+            SsoHashMap::Array(array) => array.len(),
+            SsoHashMap::Map(map) => map.len(),
+        }
+    }
+
+    /// Returns `true` if the map contains no elements.
+    pub fn is_empty(&self) -> bool {
+        match self {
+            SsoHashMap::Array(array) => array.is_empty(),
+            SsoHashMap::Map(map) => map.is_empty(),
+        }
+    }
+
+    /// An iterator visiting all key-value pairs in arbitrary order.
+    /// The iterator element type is `(&'a K, &'a V)`.
+    pub fn iter(&self) -> impl Iterator<Item = (&'_ K, &'_ V)> {
+        self.into_iter()
+    }
+
+    /// An iterator visiting all key-value pairs in arbitrary order,
+    /// with mutable references to the values.
+    /// The iterator element type is `(&'a K, &'a mut V)`.
+    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&'_ K, &'_ mut V)> {
+        self.into_iter()
+    }
+
+    /// An iterator visiting all keys in arbitrary order.
+    /// The iterator element type is `&'a K`.
+    pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.iter().map(|(k, _v)| k)),
+            SsoHashMap::Map(map) => EitherIter::Right(map.keys()),
+        }
+    }
+
+    /// An iterator visiting all values in arbitrary order.
+    /// The iterator element type is `&'a V`.
+    pub fn values(&self) -> impl Iterator<Item = &'_ V> {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.iter().map(|(_k, v)| v)),
+            SsoHashMap::Map(map) => EitherIter::Right(map.values()),
+        }
+    }
+
+    /// An iterator visiting all values mutably in arbitrary order.
+    /// The iterator element type is `&'a mut V`.
+    pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.iter_mut().map(|(_k, v)| v)),
+            SsoHashMap::Map(map) => EitherIter::Right(map.values_mut()),
+        }
+    }
+
+    /// Clears the map, returning all key-value pairs as an iterator. Keeps the
+    /// allocated memory for reuse.
+    pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.drain(..)),
+            SsoHashMap::Map(map) => EitherIter::Right(map.drain()),
+        }
+    }
+}
+
+impl<K: Eq + Hash, V> SsoHashMap<K, V> {
+    /// Changes underlying storage from array to hashmap
+    /// if array is full.
+    fn migrate_if_full(&mut self) {
+        if let SsoHashMap::Array(array) = self {
+            if array.is_full() {
+                *self = SsoHashMap::Map(array.drain(..).collect());
+            }
+        }
+    }
+
+    /// Reserves capacity for at least `additional` more elements to be inserted
+    /// in the `SsoHashMap`. The collection may reserve more space to avoid
+    /// frequent reallocations.
+    pub fn reserve(&mut self, additional: usize) {
+        match self {
+            SsoHashMap::Array(array) => {
+                if array.capacity() < (array.len() + additional) {
+                    let mut map: FxHashMap<K, V> = array.drain(..).collect();
+                    map.reserve(additional);
+                    *self = SsoHashMap::Map(map);
+                }
+            }
+            SsoHashMap::Map(map) => map.reserve(additional),
+        }
+    }
+
+    /// Shrinks the capacity of the map as much as possible. It will drop
+    /// down as much as possible while maintaining the internal rules
+    /// and possibly leaving some space in accordance with the resize policy.
+    pub fn shrink_to_fit(&mut self) {
+        if let SsoHashMap::Map(map) = self {
+            let mut array = ArrayVec::new();
+            if map.len() <= array.capacity() {
+                array.extend(map.drain());
+                *self = SsoHashMap::Array(array);
+            } else {
+                map.shrink_to_fit();
+            }
+        }
+    }
+
+    /// Retains only the elements specified by the predicate.
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        match self {
+            SsoHashMap::Array(array) => array.retain(|(k, v)| f(k, v)),
+            SsoHashMap::Map(map) => map.retain(f),
+        }
+    }
+
+    /// Inserts a key-value pair into the map.
+    ///
+    /// If the map did not have this key present, [`None`] is returned.
+    ///
+    /// If the map did have this key present, the value is updated, and the old
+    /// value is returned. The key is not updated, though; this matters for
+    /// types that can be `==` without being identical. See the [module-level
+    /// documentation] for more.
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
         match self {
             SsoHashMap::Array(array) => {
-                for pair in array.iter_mut() {
-                    if pair.0 == key {
-                        pair.1 = value;
-                        return;
+                for (k, v) in array.iter_mut() {
+                    if *k == key {
+                        let old_value = std::mem::replace(v, value);
+                        return Some(old_value);
                     }
                 }
                 if let Err(error) = array.try_push((key, value)) {
@@ -35,27 +205,325 @@ impl<K: Eq + Hash, V> SsoHashMap<K, V> {
                     map.insert(key, value);
                     *self = SsoHashMap::Map(map);
                 }
+                None
             }
-            SsoHashMap::Map(map) => {
-                map.insert(key, value);
+            SsoHashMap::Map(map) => map.insert(key, value),
+        }
+    }
+
+    /// Removes a key from the map, returning the value at the key if the key
+    /// was previously in the map.
+    pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        match self {
+            SsoHashMap::Array(array) => {
+                if let Some(index) = array.iter().position(|(k, _v)| k.borrow() == key) {
+                    Some(array.swap_remove(index).1)
+                } else {
+                    None
+                }
+            }
+            SsoHashMap::Map(map) => map.remove(key),
+        }
+    }
+
+    /// Removes a key from the map, returning the stored key and value if the
+    /// key was previously in the map.
+    pub fn remove_entry<Q: ?Sized>(&mut self, key: &Q) -> Option<(K, V)>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        match self {
+            SsoHashMap::Array(array) => {
+                if let Some(index) = array.iter().position(|(k, _v)| k.borrow() == key) {
+                    Some(array.swap_remove(index))
+                } else {
+                    None
+                }
             }
+            SsoHashMap::Map(map) => map.remove_entry(key),
         }
     }
 
-    /// Return value by key if any.
-    pub fn get(&self, key: &K) -> Option<&V> {
+    /// Returns a reference to the value corresponding to the key.
+    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
         match self {
             SsoHashMap::Array(array) => {
-                for pair in array {
-                    if pair.0 == *key {
-                        return Some(&pair.1);
+                for (k, v) in array {
+                    if k.borrow() == key {
+                        return Some(v);
                     }
                 }
-                return None;
+                None
             }
-            SsoHashMap::Map(map) => {
-                return map.get(key);
+            SsoHashMap::Map(map) => map.get(key),
+        }
+    }
+
+    /// Returns a mutable reference to the value corresponding to the key.
+    pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        match self {
+            SsoHashMap::Array(array) => {
+                for (k, v) in array {
+                    if (*k).borrow() == key {
+                        return Some(v);
+                    }
+                }
+                None
             }
+            SsoHashMap::Map(map) => map.get_mut(key),
         }
     }
+
+    /// Returns the key-value pair corresponding to the supplied key.
+    pub fn get_key_value<Q: ?Sized>(&self, key: &Q) -> Option<(&K, &V)>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        match self {
+            SsoHashMap::Array(array) => {
+                for (k, v) in array {
+                    if k.borrow() == key {
+                        return Some((k, v));
+                    }
+                }
+                None
+            }
+            SsoHashMap::Map(map) => map.get_key_value(key),
+        }
+    }
+
+    /// Returns `true` if the map contains a value for the specified key.
+    pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        match self {
+            SsoHashMap::Array(array) => array.iter().any(|(k, _v)| k.borrow() == key),
+            SsoHashMap::Map(map) => map.contains_key(key),
+        }
+    }
+
+    /// Gets the given key's corresponding entry in the map for in-place manipulation.
+    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
+        Entry { ssomap: self, key }
+    }
+}
+
+impl<K, V> Default for SsoHashMap<K, V> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<K: Eq + Hash, V> FromIterator<(K, V)> for SsoHashMap<K, V> {
+    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> SsoHashMap<K, V> {
+        let mut map: SsoHashMap<K, V> = Default::default();
+        map.extend(iter);
+        map
+    }
+}
+
+impl<K: Eq + Hash, V> Extend<(K, V)> for SsoHashMap<K, V> {
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = (K, V)>,
+    {
+        for (key, value) in iter.into_iter() {
+            self.insert(key, value);
+        }
+    }
+
+    fn extend_one(&mut self, (k, v): (K, V)) {
+        self.insert(k, v);
+    }
+
+    fn extend_reserve(&mut self, additional: usize) {
+        match self {
+            SsoHashMap::Array(array) => {
+                if array.capacity() < (array.len() + additional) {
+                    let mut map: FxHashMap<K, V> = array.drain(..).collect();
+                    map.extend_reserve(additional);
+                    *self = SsoHashMap::Map(map);
+                }
+            }
+            SsoHashMap::Map(map) => map.extend_reserve(additional),
+        }
+    }
+}
+
+impl<'a, K, V> Extend<(&'a K, &'a V)> for SsoHashMap<K, V>
+where
+    K: Eq + Hash + Copy,
+    V: Copy,
+{
+    fn extend<T: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: T) {
+        self.extend(iter.into_iter().map(|(k, v)| (k.clone(), v.clone())))
+    }
+
+    fn extend_one(&mut self, (&k, &v): (&'a K, &'a V)) {
+        self.insert(k, v);
+    }
+
+    fn extend_reserve(&mut self, additional: usize) {
+        Extend::<(K, V)>::extend_reserve(self, additional)
+    }
+}
+
+impl<K, V> IntoIterator for SsoHashMap<K, V> {
+    type IntoIter = EitherIter<
+        <ArrayVec<[(K, V); 8]> as IntoIterator>::IntoIter,
+        <FxHashMap<K, V> as IntoIterator>::IntoIter,
+    >;
+    type Item = <Self::IntoIter as Iterator>::Item;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.into_iter()),
+            SsoHashMap::Map(map) => EitherIter::Right(map.into_iter()),
+        }
+    }
+}
+
+/// adapts Item of array reference iterator to Item of hashmap reference iterator.
+fn adapt_array_ref_it<K, V>(pair: &'a (K, V)) -> (&'a K, &'a V) {
+    let (a, b) = pair;
+    (a, b)
+}
+
+/// adapts Item of array mut reference iterator to Item of hashmap mut reference iterator.
+fn adapt_array_mut_it<K, V>(pair: &'a mut (K, V)) -> (&'a K, &'a mut V) {
+    let (a, b) = pair;
+    (a, b)
+}
+
+impl<'a, K, V> IntoIterator for &'a SsoHashMap<K, V> {
+    type IntoIter = EitherIter<
+        std::iter::Map<
+            <&'a ArrayVec<[(K, V); 8]> as IntoIterator>::IntoIter,
+            fn(&'a (K, V)) -> (&'a K, &'a V),
+        >,
+        <&'a FxHashMap<K, V> as IntoIterator>::IntoIter,
+    >;
+    type Item = <Self::IntoIter as Iterator>::Item;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.into_iter().map(adapt_array_ref_it)),
+            SsoHashMap::Map(map) => EitherIter::Right(map.into_iter()),
+        }
+    }
+}
+
+impl<'a, K, V> IntoIterator for &'a mut SsoHashMap<K, V> {
+    type IntoIter = EitherIter<
+        std::iter::Map<
+            <&'a mut ArrayVec<[(K, V); 8]> as IntoIterator>::IntoIter,
+            fn(&'a mut (K, V)) -> (&'a K, &'a mut V),
+        >,
+        <&'a mut FxHashMap<K, V> as IntoIterator>::IntoIter,
+    >;
+    type Item = <Self::IntoIter as Iterator>::Item;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            SsoHashMap::Array(array) => EitherIter::Left(array.into_iter().map(adapt_array_mut_it)),
+            SsoHashMap::Map(map) => EitherIter::Right(map.into_iter()),
+        }
+    }
+}
+
+impl<K, V> fmt::Debug for SsoHashMap<K, V>
+where
+    K: fmt::Debug,
+    V: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_map().entries(self.iter()).finish()
+    }
+}
+
+impl<'a, K, Q: ?Sized, V> Index<&'a Q> for SsoHashMap<K, V>
+where
+    K: Eq + Hash + Borrow<Q>,
+    Q: Eq + Hash,
+{
+    type Output = V;
+
+    fn index(&self, key: &Q) -> &V {
+        self.get(key).expect("no entry found for key")
+    }
+}
+
+/// A view into a single entry in a map.
+pub struct Entry<'a, K, V> {
+    ssomap: &'a mut SsoHashMap<K, V>,
+    key: K,
+}
+
+impl<'a, K: Eq + Hash, V> Entry<'a, K, V> {
+    /// Provides in-place mutable access to an occupied entry before any
+    /// potential inserts into the map.
+    pub fn and_modify<F>(self, f: F) -> Self
+    where
+        F: FnOnce(&mut V),
+    {
+        if let Some(value) = self.ssomap.get_mut(&self.key) {
+            f(value);
+        }
+        self
+    }
+
+    /// Ensures a value is in the entry by inserting the default if empty, and returns
+    /// a mutable reference to the value in the entry.
+    pub fn or_insert(self, value: V) -> &'a mut V {
+        self.or_insert_with(|| value)
+    }
+
+    /// Ensures a value is in the entry by inserting the result of the default function if empty,
+    /// and returns a mutable reference to the value in the entry.
+    pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
+        self.ssomap.migrate_if_full();
+        match self.ssomap {
+            SsoHashMap::Array(array) => {
+                let key_ref = &self.key;
+                let found_index = array.iter().position(|(k, _v)| k == key_ref);
+                let index = if let Some(index) = found_index {
+                    index
+                } else {
+                    array.try_push((self.key, default())).unwrap();
+                    array.len() - 1
+                };
+                &mut array[index].1
+            }
+            SsoHashMap::Map(map) => map.entry(self.key).or_insert_with(default),
+        }
+    }
+
+    /// Returns a reference to this entry's key.
+    pub fn key(&self) -> &K {
+        &self.key
+    }
+}
+
+impl<'a, K: Eq + Hash, V: Default> Entry<'a, K, V> {
+    /// Ensures a value is in the entry by inserting the default value if empty,
+    /// and returns a mutable reference to the value in the entry.
+    pub fn or_default(self) -> &'a mut V {
+        self.or_insert_with(Default::default)
+    }
 }