Move projection_cache into the combined undo log

6 files changed, 79 insertions, 53 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 45c298ea8bf..c5e25b0a2c7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -485,7 +485,7 @@ version = "0.0.212"
 dependencies = [
  "cargo_metadata 0.9.1",
  "if_chain",
- "itertools 0.9.0",
+ "itertools 0.8.0",
  "lazy_static 1.4.0",
  "pulldown-cmark 0.7.1",
  "quine-mc_cluskey",
@@ -1630,15 +1630,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "itertools"
-version = "0.9.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "284f18f85651fe11e8a991b2adb42cb078325c996ed026d994719efcfca1d54b"
-dependencies = [
- "either",
-]
-
-[[package]]
 name = "itoa"
 version = "0.4.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2188,7 +2179,6 @@ dependencies = [
  "rustc-workspace-hack",
  "rustc_version",
  "serde",
- "serde_json",
  "shell-escape",
  "vergen",
 ]
diff --git a/src/librustc_data_structures/snapshot_map/mod.rs b/src/librustc_data_structures/snapshot_map/mod.rs
index 81b4c35aa1a..fe3c5a8afc9 100644
--- a/src/librustc_data_structures/snapshot_map/mod.rs
+++ b/src/librustc_data_structures/snapshot_map/mod.rs
@@ -1,5 +1,6 @@
 use crate::fx::FxHashMap;
 use crate::undo_log::{Rollback, Snapshots, UndoLogs, VecLog};
+use std::borrow::{Borrow, BorrowMut};
 use std::hash::Hash;
 use std::marker::PhantomData;
 use std::ops;
@@ -10,6 +11,7 @@ pub use crate::undo_log::Snapshot;
 mod tests;
 
 pub type SnapshotMapStorage<K, V> = SnapshotMap<K, V, FxHashMap<K, V>, ()>;
+pub type SnapshotMapRef<'a, K, V, L> = SnapshotMap<K, V, &'a mut FxHashMap<K, V>, &'a mut L>;
 
 pub struct SnapshotMap<K, V, M = FxHashMap<K, V>, L = VecLog<UndoLog<K, V>>> {
     map: M,
@@ -43,16 +45,16 @@ impl<K, V, M, L> SnapshotMap<K, V, M, L> {
 impl<K, V, M, L> SnapshotMap<K, V, M, L>
 where
     K: Hash + Clone + Eq,
-    M: AsMut<FxHashMap<K, V>> + AsRef<FxHashMap<K, V>>,
+    M: BorrowMut<FxHashMap<K, V>> + Borrow<FxHashMap<K, V>>,
     L: UndoLogs<UndoLog<K, V>>,
 {
     pub fn clear(&mut self) {
-        self.map.as_mut().clear();
+        self.map.borrow_mut().clear();
         self.undo_log.clear();
     }
 
     pub fn insert(&mut self, key: K, value: V) -> bool {
-        match self.map.as_mut().insert(key.clone(), value) {
+        match self.map.borrow_mut().insert(key.clone(), value) {
             None => {
                 self.undo_log.push(UndoLog::Inserted(key));
                 true
@@ -65,7 +67,7 @@ where
     }
 
     pub fn remove(&mut self, key: K) -> bool {
-        match self.map.as_mut().remove(&key) {
+        match self.map.borrow_mut().remove(&key) {
             Some(old_value) => {
                 self.undo_log.push(UndoLog::Overwrite(key, old_value));
                 true
@@ -75,7 +77,7 @@ where
     }
 
     pub fn get(&self, key: &K) -> Option<&V> {
-        self.map.as_ref().get(key)
+        self.map.borrow().get(key)
     }
 }
 
@@ -99,11 +101,21 @@ where
 impl<'k, K, V, M, L> ops::Index<&'k K> for SnapshotMap<K, V, M, L>
 where
     K: Hash + Clone + Eq,
-    M: AsRef<FxHashMap<K, V>>,
+    M: Borrow<FxHashMap<K, V>>,
 {
     type Output = V;
     fn index(&self, key: &'k K) -> &V {
-        &self.map.as_ref()[key]
+        &self.map.borrow()[key]
+    }
+}
+
+impl<K, V, M, L> Rollback<UndoLog<K, V>> for SnapshotMap<K, V, M, L>
+where
+    K: Eq + Hash,
+    M: Rollback<UndoLog<K, V>>,
+{
+    fn reverse(&mut self, undo: UndoLog<K, V>) {
+        self.map.reverse(undo)
     }
 }
 
diff --git a/src/librustc_infer/infer/mod.rs b/src/librustc_infer/infer/mod.rs
index eabb513d2d3..685dc1ab8a4 100644
--- a/src/librustc_infer/infer/mod.rs
+++ b/src/librustc_infer/infer/mod.rs
@@ -141,7 +141,7 @@ pub struct InferCtxtInner<'tcx> {
     /// Cache for projections. This cache is snapshotted along with the infcx.
     ///
     /// Public so that `traits::project` can use it.
-    pub projection_cache: traits::ProjectionCache<'tcx>,
+    pub projection_cache: traits::ProjectionCacheStorage<'tcx>,
 
     /// We instantiate `UnificationTable` with `bounds<Ty>` because the types
     /// that might instantiate a general type variable have an order,
@@ -213,6 +213,10 @@ impl<'tcx> InferCtxtInner<'tcx> {
         }
     }
 
+    pub(crate) fn projection_cache(&mut self) -> traits::ProjectionCache<'tcx, '_> {
+        self.projection_cache.with_log(&mut self.undo_log)
+    }
+
     fn type_variables(&mut self) -> type_variable::TypeVariableTable<'tcx, '_> {
         self.type_variables.with_log(&mut self.undo_log)
     }
@@ -265,6 +269,7 @@ pub(crate) enum UndoLog<'tcx> {
     FloatUnificationTable(sv::UndoLog<ut::Delegate<ty::FloatVid>>),
     RegionConstraintCollector(region_constraints::UndoLog<'tcx>),
     RegionUnificationTable(sv::UndoLog<ut::Delegate<ty::RegionVid>>),
+    ProjectionCache(traits::UndoLog<'tcx>),
 }
 
 impl<'tcx> From<region_constraints::UndoLog<'tcx>> for UndoLog<'tcx> {
@@ -327,6 +332,12 @@ impl<'tcx> From<sv::UndoLog<ut::Delegate<ty::RegionVid>>> for UndoLog<'tcx> {
     }
 }
 
+impl<'tcx> From<traits::UndoLog<'tcx>> for UndoLog<'tcx> {
+    fn from(l: traits::UndoLog<'tcx>) -> Self {
+        Self::ProjectionCache(l)
+    }
+}
+
 pub(crate) type UnificationTable<'a, 'tcx, T> =
     ut::UnificationTable<ut::InPlace<T, &'a mut ut::UnificationStorage<T>, &'a mut Logs<'tcx>>>;
 
@@ -336,6 +347,7 @@ struct RollbackView<'tcx, 'a> {
     int_unification_table: &'a mut ut::UnificationStorage<ty::IntVid>,
     float_unification_table: &'a mut ut::UnificationStorage<ty::FloatVid>,
     region_constraints: &'a mut RegionConstraintStorage<'tcx>,
+    projection_cache: &'a mut traits::ProjectionCacheStorage<'tcx>,
 }
 
 impl<'tcx> Rollback<UndoLog<'tcx>> for RollbackView<'tcx, '_> {
@@ -349,6 +361,7 @@ impl<'tcx> Rollback<UndoLog<'tcx>> for RollbackView<'tcx, '_> {
             UndoLog::RegionUnificationTable(undo) => {
                 self.region_constraints.unification_table.reverse(undo)
             }
+            UndoLog::ProjectionCache(undo) => self.projection_cache.reverse(undo),
         }
     }
 }
@@ -885,7 +898,6 @@ impl<'tcx> InferOk<'tcx, ()> {
 
 #[must_use = "once you start a snapshot, you should always consume it"]
 pub struct CombinedSnapshot<'a, 'tcx> {
-    projection_cache_snapshot: traits::ProjectionCacheSnapshot,
     undo_snapshot: Snapshot<'tcx>,
     type_snapshot: type_variable::Snapshot<'tcx>,
     const_snapshot: usize,
@@ -1016,7 +1028,6 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
         inner.undo_log.num_open_snapshots += 1;
         let undo_snapshot = Snapshot { undo_len: inner.undo_log.logs.len(), _marker: PhantomData };
         CombinedSnapshot {
-            projection_cache_snapshot: inner.projection_cache.snapshot(),
             undo_snapshot,
             type_snapshot: inner.type_variables().snapshot(),
             const_snapshot: inner.const_unification_table().len(),
@@ -1036,7 +1047,6 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
     fn rollback_to(&self, cause: &str, snapshot: CombinedSnapshot<'a, 'tcx>) {
         debug!("rollback_to(cause={})", cause);
         let CombinedSnapshot {
-            projection_cache_snapshot,
             undo_snapshot,
             type_snapshot: _,
             const_snapshot: _,
@@ -1062,6 +1072,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
             int_unification_table,
             float_unification_table,
             region_constraints,
+            projection_cache,
             ..
         } = inner;
         inner.undo_log.rollback_to(
@@ -1071,17 +1082,16 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                 int_unification_table,
                 float_unification_table,
                 region_constraints: region_constraints.as_mut().unwrap(),
+                projection_cache,
             },
             undo_snapshot,
         );
-        inner.projection_cache.rollback_to(projection_cache_snapshot);
         inner.region_obligations.truncate(region_obligations_snapshot);
     }
 
     fn commit_from(&self, snapshot: CombinedSnapshot<'a, 'tcx>) {
         debug!("commit_from()");
         let CombinedSnapshot {
-            projection_cache_snapshot,
             undo_snapshot,
             type_snapshot: _,
             const_snapshot: _,
@@ -1100,7 +1110,6 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
 
         let mut inner = self.inner.borrow_mut();
         inner.undo_log.commit(undo_snapshot);
-        inner.projection_cache.commit(projection_cache_snapshot);
     }
 
     /// Executes `f` and commit the bindings.
@@ -1773,7 +1782,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
     pub fn clear_caches(&self) {
         self.selection_cache.clear();
         self.evaluation_cache.clear();
-        self.inner.borrow_mut().projection_cache.clear();
+        self.inner.borrow_mut().projection_cache().clear();
     }
 
     fn universe(&self) -> ty::UniverseIndex {
diff --git a/src/librustc_infer/traits/mod.rs b/src/librustc_infer/traits/mod.rs
index 2210c663d14..640ea7774cd 100644
--- a/src/librustc_infer/traits/mod.rs
+++ b/src/librustc_infer/traits/mod.rs
@@ -20,6 +20,7 @@ pub use self::Vtable::*;
 
 pub use self::engine::{TraitEngine, TraitEngineExt};
 pub use self::project::MismatchedProjectionTypes;
+pub(crate) use self::project::UndoLog;
 pub use self::project::{
     Normalized, NormalizedTy, ProjectionCache, ProjectionCacheEntry, ProjectionCacheKey,
     ProjectionCacheSnapshot, Reveal,
diff --git a/src/librustc_infer/traits/project.rs b/src/librustc_infer/traits/project.rs
index 17105f99ac0..0c51dafef6f 100644
--- a/src/librustc_infer/traits/project.rs
+++ b/src/librustc_infer/traits/project.rs
@@ -8,6 +8,9 @@ use rustc_middle::ty::{self, Ty};
 
 pub use rustc_middle::traits::Reveal;
 
+pub(crate) type UndoLog<'tcx> =
+    snapshot_map::UndoLog<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>;
+
 #[derive(Clone)]
 pub struct MismatchedProjectionTypes<'tcx> {
     pub err: ty::error::TypeError<'tcx>,
@@ -58,9 +61,14 @@ impl<'tcx, T> Normalized<'tcx, T> {
 //
 // FIXME: we probably also want some sort of cross-infcx cache here to
 // reduce the amount of duplication. Let's see what we get with the Chalk reforms.
+pub struct ProjectionCache<'tcx, 'a> {
+    map: &'a mut SnapshotMapStorage<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>,
+    undo_log: &'a mut Logs<'tcx>,
+}
+
 #[derive(Default)]
-pub struct ProjectionCache<'tcx> {
-    map: SnapshotMap<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>,
+pub struct ProjectionCacheStorage<'tcx> {
+    map: SnapshotMapStorage<ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>>,
 }
 
 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
@@ -82,26 +90,24 @@ pub enum ProjectionCacheEntry<'tcx> {
     NormalizedTy(NormalizedTy<'tcx>),
 }
 
-// N.B., intentionally not Clone
-pub struct ProjectionCacheSnapshot {
-    snapshot: Snapshot,
-}
-
-impl<'tcx> ProjectionCache<'tcx> {
-    pub fn clear(&mut self) {
-        self.map.clear();
-    }
-
-    pub fn snapshot(&mut self) -> ProjectionCacheSnapshot {
-        ProjectionCacheSnapshot { snapshot: self.map.snapshot() }
+impl<'tcx> ProjectionCacheStorage<'tcx> {
+    pub(crate) fn with_log<'a>(
+        &'a mut self,
+        undo_log: &'a mut Logs<'tcx>,
+    ) -> ProjectionCache<'tcx, 'a> {
+        ProjectionCache { map: &mut self.map, undo_log }
     }
+}
 
-    pub fn rollback_to(&mut self, snapshot: ProjectionCacheSnapshot) {
-        self.map.rollback_to(snapshot.snapshot);
+impl<'tcx> ProjectionCache<'tcx, '_> {
+    fn map(
+        &mut self,
+    ) -> SnapshotMapRef<'_, ProjectionCacheKey<'tcx>, ProjectionCacheEntry<'tcx>, Logs<'tcx>> {
+        self.map.with_log(self.undo_log)
     }
 
-    pub fn commit(&mut self, snapshot: ProjectionCacheSnapshot) {
-        self.map.commit(snapshot.snapshot);
+    pub fn clear(&mut self) {
+        self.map().clear();
     }
 
     /// Try to start normalize `key`; returns an error if
@@ -111,11 +117,12 @@ impl<'tcx> ProjectionCache<'tcx> {
         &mut self,
         key: ProjectionCacheKey<'tcx>,
     ) -> Result<(), ProjectionCacheEntry<'tcx>> {
-        if let Some(entry) = self.map.get(&key) {
+        let mut map = self.map();
+        if let Some(entry) = map.get(&key) {
             return Err(entry.clone());
         }
 
-        self.map.insert(key, ProjectionCacheEntry::InProgress);
+        map.insert(key, ProjectionCacheEntry::InProgress);
         Ok(())
     }
 
@@ -125,7 +132,7 @@ impl<'tcx> ProjectionCache<'tcx> {
             "ProjectionCacheEntry::insert_ty: adding cache entry: key={:?}, value={:?}",
             key, value
         );
-        let fresh_key = self.map.insert(key, ProjectionCacheEntry::NormalizedTy(value));
+        let fresh_key = self.map().insert(key, ProjectionCacheEntry::NormalizedTy(value));
         assert!(!fresh_key, "never started projecting `{:?}`", key);
     }
 
@@ -134,7 +141,8 @@ impl<'tcx> ProjectionCache<'tcx> {
     /// snapshot - if the snapshot is rolled back, the obligations will be
     /// marked as incomplete again).
     pub fn complete(&mut self, key: ProjectionCacheKey<'tcx>) {
-        let ty = match self.map.get(&key) {
+        let mut map = self.map();
+        let ty = match map.get(&key) {
             Some(&ProjectionCacheEntry::NormalizedTy(ref ty)) => {
                 debug!("ProjectionCacheEntry::complete({:?}) - completing {:?}", key, ty);
                 ty.value
@@ -147,7 +155,7 @@ impl<'tcx> ProjectionCache<'tcx> {
             }
         };
 
-        self.map.insert(
+        map.insert(
             key,
             ProjectionCacheEntry::NormalizedTy(Normalized { value: ty, obligations: vec![] }),
         );
@@ -159,7 +167,7 @@ impl<'tcx> ProjectionCache<'tcx> {
         // We want to insert `ty` with no obligations. If the existing value
         // already has no obligations (as is common) we don't insert anything.
         if !ty.obligations.is_empty() {
-            self.map.insert(
+            self.map().insert(
                 key,
                 ProjectionCacheEntry::NormalizedTy(Normalized {
                     value: ty.value,
@@ -174,14 +182,20 @@ impl<'tcx> ProjectionCache<'tcx> {
     /// type information (in which case, the "fully resolved" key will
     /// be different).
     pub fn ambiguous(&mut self, key: ProjectionCacheKey<'tcx>) {
-        let fresh = self.map.insert(key, ProjectionCacheEntry::Ambiguous);
+        let fresh = self.map().insert(key, ProjectionCacheEntry::Ambiguous);
         assert!(!fresh, "never started projecting `{:?}`", key);
     }
 
     /// Indicates that trying to normalize `key` resulted in
     /// error.
     pub fn error(&mut self, key: ProjectionCacheKey<'tcx>) {
-        let fresh = self.map.insert(key, ProjectionCacheEntry::Error);
+        let fresh = self.map().insert(key, ProjectionCacheEntry::Error);
         assert!(!fresh, "never started projecting `{:?}`", key);
     }
 }
+
+impl<'tcx> Rollback<UndoLog<'tcx>> for ProjectionCacheStorage<'tcx> {
+    fn reverse(&mut self, undo: UndoLog<'tcx>) {
+        self.map.reverse(undo);
+    }
+}
diff --git a/src/librustc_trait_selection/traits/select.rs b/src/librustc_trait_selection/traits/select.rs
index dfbb0742448..38590abf1f7 100644
--- a/src/librustc_trait_selection/traits/select.rs
+++ b/src/librustc_trait_selection/traits/select.rs
@@ -471,7 +471,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                         if let Some(key) =
                             ProjectionCacheKey::from_poly_projection_predicate(self, data)
                         {
-                            self.infcx.inner.borrow_mut().projection_cache.complete(key);
+                            self.infcx.inner.borrow_mut().projection_cache().complete(key);
                         }
                         result
                     }