SsoHashSet reimplemented as a wrapper on top of SsoHashMap

SsoHashSet::replace had to be removed because
it requires missing API from SsoHashMap.
It's not a widely used function, so I think it's ok
to omit it for now.

EitherIter moved into its own file.

Also sprinkled code with #[inline] attributes where appropriate.

4 files changed, 158 insertions, 228 deletions

diff --git a/compiler/rustc_data_structures/src/sso/either_iter.rs b/compiler/rustc_data_structures/src/sso/either_iter.rs
new file mode 100644
index 00000000000..af8ffcf4c13
--- /dev/null
+++ b/compiler/rustc_data_structures/src/sso/either_iter.rs
@@ -0,0 +1,75 @@
+use std::fmt;
+use std::iter::ExactSizeIterator;
+use std::iter::FusedIterator;
+use std::iter::Iterator;
+
+/// Iterator which may contain instance of
+/// one of two specific implementations.
+///
+/// Note: For most methods providing custom
+///       implementation may margianlly
+///       improve performance by avoiding
+///       doing Left/Right match on every step
+///       and doing it only once instead.
+#[derive(Clone)]
+pub enum EitherIter<L, R> {
+    Left(L),
+    Right(R),
+}
+
+impl<L, R> Iterator for EitherIter<L, R>
+where
+    L: Iterator,
+    R: Iterator<Item = L::Item>,
+{
+    type Item = L::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            EitherIter::Left(l) => l.next(),
+            EitherIter::Right(r) => r.next(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match self {
+            EitherIter::Left(l) => l.size_hint(),
+            EitherIter::Right(r) => r.size_hint(),
+        }
+    }
+}
+
+impl<L, R> ExactSizeIterator for EitherIter<L, R>
+where
+    L: ExactSizeIterator,
+    R: ExactSizeIterator,
+    EitherIter<L, R>: Iterator,
+{
+    fn len(&self) -> usize {
+        match self {
+            EitherIter::Left(l) => l.len(),
+            EitherIter::Right(r) => r.len(),
+        }
+    }
+}
+
+impl<L, R> FusedIterator for EitherIter<L, R>
+where
+    L: FusedIterator,
+    R: FusedIterator,
+    EitherIter<L, R>: Iterator,
+{
+}
+
+impl<L, R> fmt::Debug for EitherIter<L, R>
+where
+    L: fmt::Debug,
+    R: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            EitherIter::Left(l) => l.fmt(f),
+            EitherIter::Right(r) => r.fmt(f),
+        }
+    }
+}
diff --git a/compiler/rustc_data_structures/src/sso/map.rs b/compiler/rustc_data_structures/src/sso/map.rs
index 258368c8ef3..3089f887845 100644
--- a/compiler/rustc_data_structures/src/sso/map.rs
+++ b/compiler/rustc_data_structures/src/sso/map.rs
@@ -1,4 +1,4 @@
-use super::EitherIter;
+use super::either_iter::EitherIter;
 use crate::fx::FxHashMap;
 use arrayvec::ArrayVec;
 use std::borrow::Borrow;
@@ -32,6 +32,7 @@ pub enum SsoHashMap<K, V> {
 
 impl<K, V> SsoHashMap<K, V> {
     /// Creates an empty `SsoHashMap`.
+    #[inline]
     pub fn new() -> Self {
         SsoHashMap::Array(ArrayVec::new())
     }
@@ -81,13 +82,15 @@ impl<K, V> SsoHashMap<K, V> {
 
     /// 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)> {
+    #[inline]
+    pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
         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)`.
+    #[inline]
     pub fn iter_mut(&mut self) -> impl Iterator<Item = (&'_ K, &'_ mut V)> {
         self.into_iter()
     }
@@ -319,12 +322,14 @@ impl<K: Eq + Hash, V> SsoHashMap<K, V> {
     }
 
     /// Gets the given key's corresponding entry in the map for in-place manipulation.
+    #[inline]
     pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
         Entry { ssomap: self, key }
     }
 }
 
 impl<K, V> Default for SsoHashMap<K, V> {
+    #[inline]
     fn default() -> Self {
         Self::new()
     }
@@ -348,6 +353,7 @@ impl<K: Eq + Hash, V> Extend<(K, V)> for SsoHashMap<K, V> {
         }
     }
 
+    #[inline]
     fn extend_one(&mut self, (k, v): (K, V)) {
         self.insert(k, v);
     }
@@ -375,10 +381,12 @@ where
         self.extend(iter.into_iter().map(|(k, v)| (k.clone(), v.clone())))
     }
 
+    #[inline]
     fn extend_one(&mut self, (&k, &v): (&'a K, &'a V)) {
         self.insert(k, v);
     }
 
+    #[inline]
     fn extend_reserve(&mut self, additional: usize) {
         Extend::<(K, V)>::extend_reserve(self, additional)
     }
@@ -400,12 +408,14 @@ impl<K, V> IntoIterator for SsoHashMap<K, V> {
 }
 
 /// adapts Item of array reference iterator to Item of hashmap reference iterator.
+#[inline(always)]
 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.
+#[inline(always)]
 fn adapt_array_mut_it<K, V>(pair: &'a mut (K, V)) -> (&'a K, &'a mut V) {
     let (a, b) = pair;
     (a, b)
@@ -464,6 +474,7 @@ where
 {
     type Output = V;
 
+    #[inline]
     fn index(&self, key: &Q) -> &V {
         self.get(key).expect("no entry found for key")
     }
@@ -490,6 +501,7 @@ impl<'a, K: Eq + Hash, V> Entry<'a, K, V> {
 
     /// Ensures a value is in the entry by inserting the default if empty, and returns
     /// a mutable reference to the value in the entry.
+    #[inline]
     pub fn or_insert(self, value: V) -> &'a mut V {
         self.or_insert_with(|| value)
     }
@@ -515,6 +527,7 @@ impl<'a, K: Eq + Hash, V> Entry<'a, K, V> {
     }
 
     /// Returns a reference to this entry's key.
+    #[inline]
     pub fn key(&self) -> &K {
         &self.key
     }
@@ -523,6 +536,7 @@ impl<'a, K: Eq + Hash, V> Entry<'a, K, V> {
 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.
+    #[inline]
     pub fn or_default(self) -> &'a mut V {
         self.or_insert_with(Default::default)
     }
diff --git a/compiler/rustc_data_structures/src/sso/mod.rs b/compiler/rustc_data_structures/src/sso/mod.rs
index bcc4240721e..dd21bc8e696 100644
--- a/compiler/rustc_data_structures/src/sso/mod.rs
+++ b/compiler/rustc_data_structures/src/sso/mod.rs
@@ -1,75 +1,4 @@
-use std::fmt;
-use std::iter::ExactSizeIterator;
-use std::iter::FusedIterator;
-use std::iter::Iterator;
-
-/// Iterator which may contain instance of
-/// one of two specific implementations.
-///
-/// Used by both SsoHashMap and SsoHashSet.
-#[derive(Clone)]
-pub enum EitherIter<L, R> {
-    Left(L),
-    Right(R),
-}
-
-impl<L, R> Iterator for EitherIter<L, R>
-where
-    L: Iterator,
-    R: Iterator<Item = L::Item>,
-{
-    type Item = L::Item;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match self {
-            EitherIter::Left(l) => l.next(),
-            EitherIter::Right(r) => r.next(),
-        }
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        match self {
-            EitherIter::Left(l) => l.size_hint(),
-            EitherIter::Right(r) => r.size_hint(),
-        }
-    }
-}
-
-impl<L, R> ExactSizeIterator for EitherIter<L, R>
-where
-    L: ExactSizeIterator,
-    R: ExactSizeIterator,
-    EitherIter<L, R>: Iterator,
-{
-    fn len(&self) -> usize {
-        match self {
-            EitherIter::Left(l) => l.len(),
-            EitherIter::Right(r) => r.len(),
-        }
-    }
-}
-
-impl<L, R> FusedIterator for EitherIter<L, R>
-where
-    L: FusedIterator,
-    R: FusedIterator,
-    EitherIter<L, R>: Iterator,
-{
-}
-
-impl<L, R> fmt::Debug for EitherIter<L, R>
-where
-    L: fmt::Debug,
-    R: fmt::Debug,
-{
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            EitherIter::Left(l) => l.fmt(f),
-            EitherIter::Right(r) => r.fmt(f),
-        }
-    }
-}
-
+mod either_iter;
 mod map;
 mod set;
 
diff --git a/compiler/rustc_data_structures/src/sso/set.rs b/compiler/rustc_data_structures/src/sso/set.rs
index 6ec70524368..353529b0598 100644
--- a/compiler/rustc_data_structures/src/sso/set.rs
+++ b/compiler/rustc_data_structures/src/sso/set.rs
@@ -1,11 +1,10 @@
-use super::EitherIter;
-use crate::fx::FxHashSet;
-use arrayvec::ArrayVec;
 use std::borrow::Borrow;
 use std::fmt;
 use std::hash::Hash;
 use std::iter::FromIterator;
 
+use super::map::SsoHashMap;
+
 /// Small-storage-optimized implementation of a set.
 ///
 /// Stores elements in a small array up to a certain length
@@ -18,77 +17,73 @@ use std::iter::FromIterator;
 ///   try_reserve (unstable)
 ///   shrink_to (unstable)
 ///   drain_filter (unstable)
+///   replace
 ///   get_or_insert/get_or_insert_owned/get_or_insert_with (unstable)
 ///   difference/symmetric_difference/intersection/union
 ///   is_disjoint/is_subset/is_superset
-///   PartialEq/Eq (requires sorting the array)
+///   PartialEq/Eq (requires SsoHashMap implementation)
 ///   BitOr/BitAnd/BitXor/Sub
 #[derive(Clone)]
-pub enum SsoHashSet<T> {
-    Array(ArrayVec<[T; 8]>),
-    Set(FxHashSet<T>),
+pub struct SsoHashSet<T> {
+    map: SsoHashMap<T, ()>,
+}
+
+/// Adapter function used ot return
+/// result if SsoHashMap functions into
+/// result SsoHashSet should return.
+#[inline(always)]
+fn entry_to_key<K, V>((k, _v): (K, V)) -> K {
+    k
 }
 
 impl<T> SsoHashSet<T> {
     /// Creates an empty `SsoHashSet`.
+    #[inline]
     pub fn new() -> Self {
-        SsoHashSet::Array(ArrayVec::new())
+        Self { map: SsoHashMap::new() }
     }
 
     /// Creates an empty `SsoHashSet` with the specified capacity.
+    #[inline]
     pub fn with_capacity(cap: usize) -> Self {
-        let array = ArrayVec::new();
-        if array.capacity() >= cap {
-            SsoHashSet::Array(array)
-        } else {
-            SsoHashSet::Set(FxHashSet::with_capacity_and_hasher(cap, Default::default()))
-        }
+        Self { map: SsoHashMap::with_capacity(cap) }
     }
 
     /// Clears the set, removing all values.
+    #[inline]
     pub fn clear(&mut self) {
-        match self {
-            SsoHashSet::Array(array) => array.clear(),
-            SsoHashSet::Set(set) => set.clear(),
-        }
+        self.map.clear()
     }
 
     /// Returns the number of elements the set can hold without reallocating.
+    #[inline]
     pub fn capacity(&self) -> usize {
-        match self {
-            SsoHashSet::Array(array) => array.capacity(),
-            SsoHashSet::Set(set) => set.capacity(),
-        }
+        self.map.capacity()
     }
 
     /// Returns the number of elements in the set.
+    #[inline]
     pub fn len(&self) -> usize {
-        match self {
-            SsoHashSet::Array(array) => array.len(),
-            SsoHashSet::Set(set) => set.len(),
-        }
+        self.map.len()
     }
 
     /// Returns `true` if the set contains no elements.
+    #[inline]
     pub fn is_empty(&self) -> bool {
-        match self {
-            SsoHashSet::Array(array) => array.is_empty(),
-            SsoHashSet::Set(set) => set.is_empty(),
-        }
+        self.map.is_empty()
     }
 
     /// An iterator visiting all elements in arbitrary order.
     /// The iterator element type is `&'a T`.
+    #[inline]
     pub fn iter(&'a self) -> impl Iterator<Item = &'a T> {
         self.into_iter()
     }
 
     /// Clears the set, returning all elements in an iterator.
+    #[inline]
     pub fn drain(&mut self) -> impl Iterator<Item = T> + '_ {
-        match self {
-            SsoHashSet::Array(array) => EitherIter::Left(array.drain(..)),
-            SsoHashSet::Set(set) => EitherIter::Right(set.drain()),
-        }
+        self.map.drain().map(entry_to_key)
     }
 }
 
@@ -96,95 +91,46 @@ impl<T: Eq + Hash> SsoHashSet<T> {
     /// Reserves capacity for at least `additional` more elements to be inserted
     /// in the `SsoHashSet`. The collection may reserve more space to avoid
     /// frequent reallocations.
+    #[inline]
     pub fn reserve(&mut self, additional: usize) {
-        match self {
-            SsoHashSet::Array(array) => {
-                if array.capacity() < (array.len() + additional) {
-                    let mut set: FxHashSet<T> = array.drain(..).collect();
-                    set.reserve(additional);
-                    *self = SsoHashSet::Set(set);
-                }
-            }
-            SsoHashSet::Set(set) => set.reserve(additional),
-        }
+        self.map.reserve(additional)
     }
 
     /// Shrinks the capacity of the set 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.
+    #[inline]
     pub fn shrink_to_fit(&mut self) {
-        if let SsoHashSet::Set(set) = self {
-            let mut array = ArrayVec::new();
-            if set.len() <= array.capacity() {
-                array.extend(set.drain());
-                *self = SsoHashSet::Array(array);
-            } else {
-                set.shrink_to_fit();
-            }
-        }
+        self.map.shrink_to_fit()
     }
 
     /// Retains only the elements specified by the predicate.
+    #[inline]
     pub fn retain<F>(&mut self, mut f: F)
     where
         F: FnMut(&T) -> bool,
     {
-        match self {
-            SsoHashSet::Array(array) => array.retain(|v| f(v)),
-            SsoHashSet::Set(set) => set.retain(f),
-        }
+        self.map.retain(|k, _v| f(k))
     }
 
     /// Removes and returns the value in the set, if any, that is equal to the given one.
+    #[inline]
     pub fn take<Q: ?Sized>(&mut self, value: &Q) -> Option<T>
     where
         T: Borrow<Q>,
         Q: Hash + Eq,
     {
-        match self {
-            SsoHashSet::Array(array) => {
-                if let Some(index) = array.iter().position(|val| val.borrow() == value) {
-                    Some(array.swap_remove(index))
-                } else {
-                    None
-                }
-            }
-            SsoHashSet::Set(set) => set.take(value),
-        }
-    }
-
-    /// Adds a value to the set, replacing the existing value, if any, that is equal to the given
-    /// one. Returns the replaced value.
-    pub fn replace(&mut self, value: T) -> Option<T> {
-        match self {
-            SsoHashSet::Array(array) => {
-                if let Some(index) = array.iter().position(|val| *val == value) {
-                    let old_value = std::mem::replace(&mut array[index], value);
-                    Some(old_value)
-                } else {
-                    None
-                }
-            }
-            SsoHashSet::Set(set) => set.replace(value),
-        }
+        self.map.remove_entry(value).map(entry_to_key)
     }
 
     /// Returns a reference to the value in the set, if any, that is equal to the given value.
+    #[inline]
     pub fn get<Q: ?Sized>(&self, value: &Q) -> Option<&T>
     where
         T: Borrow<Q>,
         Q: Hash + Eq,
     {
-        match self {
-            SsoHashSet::Array(array) => {
-                if let Some(index) = array.iter().position(|val| val.borrow() == value) {
-                    Some(&array[index])
-                } else {
-                    None
-                }
-            }
-            SsoHashSet::Set(set) => set.get(value),
-        }
+        self.map.get_key_value(value).map(entry_to_key)
     }
 
     /// Adds a value to the set.
@@ -192,60 +138,30 @@ impl<T: Eq + Hash> SsoHashSet<T> {
     /// If the set did not have this value present, `true` is returned.
     ///
     /// If the set did have this value present, `false` is returned.
+    #[inline]
     pub fn insert(&mut self, elem: T) -> bool {
-        match self {
-            SsoHashSet::Array(array) => {
-                if array.iter().any(|e| *e == elem) {
-                    false
-                } else {
-                    if let Err(error) = array.try_push(elem) {
-                        let mut set: FxHashSet<T> = array.drain(..).collect();
-                        set.insert(error.element());
-                        *self = SsoHashSet::Set(set);
-                    }
-                    true
-                }
-            }
-            SsoHashSet::Set(set) => set.insert(elem),
-        }
+        self.map.insert(elem, ()).is_none()
     }
 
     /// Removes a value from the set. Returns whether the value was
     /// present in the set.
+    #[inline]
     pub fn remove<Q: ?Sized>(&mut self, value: &Q) -> bool
     where
         T: Borrow<Q>,
         Q: Hash + Eq,
     {
-        match self {
-            SsoHashSet::Array(array) => {
-                if let Some(index) = array.iter().position(|val| val.borrow() == value) {
-                    array.swap_remove(index);
-                    true
-                } else {
-                    false
-                }
-            }
-            SsoHashSet::Set(set) => set.remove(value),
-        }
+        self.map.remove(value).is_some()
     }
 
     /// Returns `true` if the set contains a value.
+    #[inline]
     pub fn contains<Q: ?Sized>(&self, value: &Q) -> bool
     where
         T: Borrow<Q>,
         Q: Hash + Eq,
     {
-        match self {
-            SsoHashSet::Array(array) => array.iter().any(|v| v.borrow() == value),
-            SsoHashSet::Set(set) => set.contains(value),
-        }
-    }
-}
-
-impl<T> Default for SsoHashSet<T> {
-    fn default() -> Self {
-        Self::new()
+        self.map.contains_key(value)
     }
 }
 
@@ -257,6 +173,13 @@ impl<T: Eq + Hash> FromIterator<T> for SsoHashSet<T> {
     }
 }
 
+impl<T> Default for SsoHashSet<T> {
+    #[inline]
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl<T: Eq + Hash> Extend<T> for SsoHashSet<T> {
     fn extend<I>(&mut self, iter: I)
     where
@@ -267,21 +190,14 @@ impl<T: Eq + Hash> Extend<T> for SsoHashSet<T> {
         }
     }
 
+    #[inline]
     fn extend_one(&mut self, item: T) {
         self.insert(item);
     }
 
+    #[inline]
     fn extend_reserve(&mut self, additional: usize) {
-        match self {
-            SsoHashSet::Array(array) => {
-                if array.capacity() < (array.len() + additional) {
-                    let mut set: FxHashSet<T> = array.drain(..).collect();
-                    set.extend_reserve(additional);
-                    *self = SsoHashSet::Set(set);
-                }
-            }
-            SsoHashSet::Set(set) => set.extend_reserve(additional),
-        }
+        self.map.extend_reserve(additional)
     }
 }
 
@@ -289,46 +205,42 @@ impl<'a, T> Extend<&'a T> for SsoHashSet<T>
 where
     T: 'a + Eq + Hash + Copy,
 {
+    #[inline]
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
         self.extend(iter.into_iter().cloned());
     }
 
+    #[inline]
     fn extend_one(&mut self, &item: &'a T) {
         self.insert(item);
     }
 
+    #[inline]
     fn extend_reserve(&mut self, additional: usize) {
         Extend::<T>::extend_reserve(self, additional)
     }
 }
 
 impl<T> IntoIterator for SsoHashSet<T> {
-    type IntoIter = EitherIter<
-        <ArrayVec<[T; 8]> as IntoIterator>::IntoIter,
-        <FxHashSet<T> as IntoIterator>::IntoIter,
-    >;
+    type IntoIter = std::iter::Map<<SsoHashMap<T, ()> as IntoIterator>::IntoIter, fn((T, ())) -> T>;
     type Item = <Self::IntoIter as Iterator>::Item;
 
+    #[inline]
     fn into_iter(self) -> Self::IntoIter {
-        match self {
-            SsoHashSet::Array(array) => EitherIter::Left(array.into_iter()),
-            SsoHashSet::Set(set) => EitherIter::Right(set.into_iter()),
-        }
+        self.map.into_iter().map(entry_to_key)
     }
 }
 
 impl<'a, T> IntoIterator for &'a SsoHashSet<T> {
-    type IntoIter = EitherIter<
-        <&'a ArrayVec<[T; 8]> as IntoIterator>::IntoIter,
-        <&'a FxHashSet<T> as IntoIterator>::IntoIter,
+    type IntoIter = std::iter::Map<
+        <&'a SsoHashMap<T, ()> as IntoIterator>::IntoIter,
+        fn((&'a T, &'a ())) -> &'a T,
     >;
     type Item = <Self::IntoIter as Iterator>::Item;
 
+    #[inline]
     fn into_iter(self) -> Self::IntoIter {
-        match self {
-            SsoHashSet::Array(array) => EitherIter::Left(array.into_iter()),
-            SsoHashSet::Set(set) => EitherIter::Right(set.into_iter()),
-        }
+        self.map.iter().map(entry_to_key)
     }
 }