cache types during normalization

3 files changed, 69 insertions, 5 deletions

diff --git a/compiler/rustc_data_structures/Cargo.toml b/compiler/rustc_data_structures/Cargo.toml
index 3dc55cab95a..7a6f4b82b52 100644
--- a/compiler/rustc_data_structures/Cargo.toml
+++ b/compiler/rustc_data_structures/Cargo.toml
@@ -8,6 +8,7 @@ edition = "2018"
 doctest = false
 
 [dependencies]
+arrayvec = { version = "0.5.1", default-features = false }
 ena = "0.14"
 indexmap = "1.5.1"
 tracing = "0.1"
diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 88c160e93b6..7d4aed14be1 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -87,25 +87,27 @@ pub mod sorted_map;
 pub mod stable_set;
 #[macro_use]
 pub mod stable_hasher;
+mod atomic_ref;
+pub mod fingerprint;
+pub mod profiling;
 pub mod sharded;
 pub mod stack;
 pub mod sync;
 pub mod thin_vec;
 pub mod tiny_list;
 pub mod transitive_relation;
-pub use ena::undo_log;
-pub use ena::unify;
-mod atomic_ref;
-pub mod fingerprint;
-pub mod profiling;
 pub mod vec_linked_list;
 pub mod work_queue;
 pub use atomic_ref::AtomicRef;
 pub mod frozen;
+pub mod mini_map;
 pub mod tagged_ptr;
 pub mod temp_dir;
 pub mod unhash;
 
+pub use ena::undo_log;
+pub use ena::unify;
+
 pub struct OnDrop<F: Fn()>(pub F);
 
 impl<F: Fn()> OnDrop<F> {
diff --git a/compiler/rustc_data_structures/src/mini_map.rs b/compiler/rustc_data_structures/src/mini_map.rs
new file mode 100644
index 00000000000..cd3e949d383
--- /dev/null
+++ b/compiler/rustc_data_structures/src/mini_map.rs
@@ -0,0 +1,61 @@
+use crate::fx::FxHashMap;
+use arrayvec::ArrayVec;
+
+use std::hash::Hash;
+
+/// Small-storage-optimized implementation of a map
+/// made specifically for caching results.
+///
+/// Stores elements in a small array up to a certain length
+/// and switches to `HashMap` when that length is exceeded.
+pub enum MiniMap<K, V> {
+    Array(ArrayVec<[(K, V); 8]>),
+    Map(FxHashMap<K, V>),
+}
+
+impl<K: Eq + Hash, V> MiniMap<K, V> {
+    /// Creates an empty `MiniMap`.
+    pub fn new() -> Self {
+        MiniMap::Array(ArrayVec::new())
+    }
+
+    /// Inserts or updates value in the map.
+    pub fn insert(&mut self, key: K, value: V) {
+        match self {
+            MiniMap::Array(array) => {
+                for pair in array.iter_mut() {
+                    if pair.0 == key {
+                        pair.1 = value;
+                        return;
+                    }
+                }
+                if let Err(error) = array.try_push((key, value)) {
+                    let mut map: FxHashMap<K, V> = array.drain(..).collect();
+                    let (key, value) = error.element();
+                    map.insert(key, value);
+                    *self = MiniMap::Map(map);
+                }
+            }
+            MiniMap::Map(map) => {
+                map.insert(key, value);
+            }
+        }
+    }
+
+    /// Return value by key if any.
+    pub fn get(&self, key: &K) -> Option<&V> {
+        match self {
+            MiniMap::Array(array) => {
+                for pair in array {
+                    if pair.0 == *key {
+                        return Some(&pair.1);
+                    }
+                }
+                return None;
+            }
+            MiniMap::Map(map) => {
+                return map.get(key);
+            }
+        }
+    }
+}