Add roll back infrastructure for opaque type caches

10 files changed, 113 insertions, 37 deletions

diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index 73103643e3e..75305d9c923 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -190,7 +190,8 @@ pub(crate) fn type_check<'mir, 'tcx>(
             liveness::generate(&mut cx, body, elements, flow_inits, move_data, location_table);
 
             translate_outlives_facts(&mut cx);
-            let opaque_type_values = mem::take(&mut infcx.inner.borrow_mut().opaque_types);
+            let opaque_type_values =
+                infcx.inner.borrow_mut().opaque_type_storage.take_opaque_types();
 
             opaque_type_values
                 .into_iter()
@@ -1315,8 +1316,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         // have to solve any bounds (e.g., `-> impl Iterator` needs to
         // prove that `T: Iterator` where `T` is the type we
         // instantiated it with).
-        let opaque_type_map = self.infcx.inner.borrow().opaque_types.clone();
-        for (opaque_type_key, opaque_decl) in opaque_type_map {
+        for (opaque_type_key, opaque_decl) in self.infcx.opaque_types() {
             self.fully_perform_op(
                 locations,
                 ConstraintCategory::OpaqueType,
diff --git a/compiler/rustc_data_structures/src/vec_map.rs b/compiler/rustc_data_structures/src/vec_map.rs
index cc7ec9432fa..8f525230e7e 100644
--- a/compiler/rustc_data_structures/src/vec_map.rs
+++ b/compiler/rustc_data_structures/src/vec_map.rs
@@ -30,6 +30,11 @@ where
         }
     }
 
+    /// Removes the entry from the map and returns the removed value
+    pub fn remove(&mut self, k: &K) -> Option<V> {
+        self.0.iter().position(|(k2, _)| k2 == k).map(|pos| self.0.remove(pos).1)
+    }
+
     /// Gets a reference to the value in the entry.
     pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
     where
diff --git a/compiler/rustc_infer/src/infer/equate.rs b/compiler/rustc_infer/src/infer/equate.rs
index 90c0ff9226f..06692be4f72 100644
--- a/compiler/rustc_infer/src/infer/equate.rs
+++ b/compiler/rustc_infer/src/infer/equate.rs
@@ -66,17 +66,18 @@ impl<'tcx> TypeRelation<'tcx> for Equate<'_, '_, 'tcx> {
         self.relate(a, b)
     }
 
+    #[instrument(skip(self), level = "debug")]
     fn tys(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
-        debug!("{}.tys({:?}, {:?})", self.tag(), a, b);
         if a == b {
             return Ok(a);
         }
 
+        trace!(a = ?a.kind(), b = ?b.kind());
+
         let infcx = self.fields.infcx;
         let a = infcx.inner.borrow_mut().type_variables().replace_if_possible(a);
         let b = infcx.inner.borrow_mut().type_variables().replace_if_possible(b);
-
-        debug!("{}.tys: replacements ({:?}, {:?})", self.tag(), a, b);
+        trace!(a = ?a.kind(), b = ?b.kind(), "replacements");
 
         match (a.kind(), b.kind()) {
             (&ty::Infer(TyVar(a_id)), &ty::Infer(TyVar(b_id))) => {
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index d1b24b332bd..6c7e079a724 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -5,7 +5,7 @@ pub use self::RegionVariableOrigin::*;
 pub use self::SubregionOrigin::*;
 pub use self::ValuePairs::*;
 
-use self::opaque_types::OpaqueTypeMap;
+use self::opaque_types::OpaqueTypeStorage;
 pub(crate) use self::undo_log::{InferCtxtUndoLogs, Snapshot, UndoLog};
 
 use crate::traits::{self, ObligationCause, PredicateObligations, TraitEngine};
@@ -190,20 +190,10 @@ pub struct InferCtxtInner<'tcx> {
     /// that all type inference variables have been bound and so forth.
     region_obligations: Vec<(hir::HirId, RegionObligation<'tcx>)>,
 
-    undo_log: InferCtxtUndoLogs<'tcx>,
+    /// Caches for opaque type inference.
+    pub opaque_type_storage: OpaqueTypeStorage<'tcx>,
 
-    // Opaque types found in explicit return types and their
-    // associated fresh inference variable. Writeback resolves these
-    // variables to get the concrete type, which can be used to
-    // 'de-opaque' OpaqueTypeDecl outside of type inference.
-    pub opaque_types: OpaqueTypeMap<'tcx>,
-
-    /// A map from inference variables created from opaque
-    /// type instantiations (`ty::Infer`) to the actual opaque
-    /// type (`ty::Opaque`). Used during fallback to map unconstrained
-    /// opaque type inference variables to their corresponding
-    /// opaque type.
-    pub opaque_types_vars: FxHashMap<Ty<'tcx>, Ty<'tcx>>,
+    undo_log: InferCtxtUndoLogs<'tcx>,
 }
 
 impl<'tcx> InferCtxtInner<'tcx> {
@@ -217,8 +207,7 @@ impl<'tcx> InferCtxtInner<'tcx> {
             float_unification_storage: ut::UnificationTableStorage::new(),
             region_constraint_storage: Some(RegionConstraintStorage::new()),
             region_obligations: vec![],
-            opaque_types: Default::default(),
-            opaque_types_vars: Default::default(),
+            opaque_type_storage: Default::default(),
         }
     }
 
@@ -238,6 +227,11 @@ impl<'tcx> InferCtxtInner<'tcx> {
     }
 
     #[inline]
+    fn opaque_types(&mut self) -> opaque_types::OpaqueTypeTable<'_, 'tcx> {
+        self.opaque_type_storage.with_log(&mut self.undo_log)
+    }
+
+    #[inline]
     fn int_unification_table(
         &mut self,
     ) -> ut::UnificationTable<
diff --git a/compiler/rustc_infer/src/infer/opaque_types.rs b/compiler/rustc_infer/src/infer/opaque_types.rs
index e7dca94806c..e3e48a7f890 100644
--- a/compiler/rustc_infer/src/infer/opaque_types.rs
+++ b/compiler/rustc_infer/src/infer/opaque_types.rs
@@ -14,6 +14,10 @@ use std::ops::ControlFlow;
 
 pub type OpaqueTypeMap<'tcx> = VecMap<OpaqueTypeKey<'tcx>, OpaqueTypeDecl<'tcx>>;
 
+mod table;
+
+pub use table::{OpaqueTypeStorage, OpaqueTypeTable};
+
 /// Information about the opaque types whose values we
 /// are inferring in this function (these are the `impl Trait` that
 /// appear in the return type).
@@ -352,6 +356,10 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
         };
         in_definition_scope.then_some(*origin)
     }
+
+    pub fn opaque_types(&self) -> OpaqueTypeMap<'tcx> {
+        self.inner.borrow().opaque_type_storage.opaque_types()
+    }
 }
 
 // Visitor that requires that (almost) all regions in the type visited outlive
@@ -513,7 +521,9 @@ impl<'a, 'tcx> Instantiator<'a, 'tcx> {
 
         // Use the same type variable if the exact same opaque type appears more
         // than once in the return type (e.g., if it's passed to a type alias).
-        if let Some(opaque_defn) = infcx.inner.borrow().opaque_types.get(&opaque_type_key) {
+        if let Some(opaque_defn) =
+            infcx.inner.borrow().opaque_type_storage.get_decl(&opaque_type_key)
+        {
             debug!("re-using cached concrete type {:?}", opaque_defn.concrete_ty.kind());
             return opaque_defn.concrete_ty;
         }
@@ -530,14 +540,10 @@ impl<'a, 'tcx> Instantiator<'a, 'tcx> {
         // Foo, impl Bar)`.
         let definition_span = self.value_span;
 
-        {
-            let mut infcx = self.infcx.inner.borrow_mut();
-            infcx.opaque_types.insert(
-                OpaqueTypeKey { def_id, substs },
-                OpaqueTypeDecl { opaque_type: ty, definition_span, concrete_ty: ty_var, origin },
-            );
-            infcx.opaque_types_vars.insert(ty_var, ty);
-        }
+        self.infcx.inner.borrow_mut().opaque_types().register(
+            OpaqueTypeKey { def_id, substs },
+            OpaqueTypeDecl { opaque_type: ty, definition_span, concrete_ty: ty_var, origin },
+        );
 
         debug!("generated new type inference var {:?}", ty_var.kind());
 
diff --git a/compiler/rustc_infer/src/infer/opaque_types/table.rs b/compiler/rustc_infer/src/infer/opaque_types/table.rs
new file mode 100644
index 00000000000..11eeeb08c98
--- /dev/null
+++ b/compiler/rustc_infer/src/infer/opaque_types/table.rs
@@ -0,0 +1,69 @@
+use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::undo_log::UndoLogs;
+use rustc_middle::ty::{OpaqueTypeKey, Ty};
+
+use crate::infer::InferCtxtUndoLogs;
+
+use super::{OpaqueTypeDecl, OpaqueTypeMap};
+
+#[derive(Default)]
+pub struct OpaqueTypeStorage<'tcx> {
+    // Opaque types found in explicit return types and their
+    // associated fresh inference variable. Writeback resolves these
+    // variables to get the concrete type, which can be used to
+    // 'de-opaque' OpaqueTypeDecl, after typeck is done with all functions.
+    pub opaque_types: OpaqueTypeMap<'tcx>,
+
+    /// A map from inference variables created from opaque
+    /// type instantiations (`ty::Infer`) to the actual opaque
+    /// type (`ty::Opaque`). Used during fallback to map unconstrained
+    /// opaque type inference variables to their corresponding
+    /// opaque type.
+    pub opaque_types_vars: FxHashMap<Ty<'tcx>, Ty<'tcx>>,
+}
+
+impl<'tcx> OpaqueTypeStorage<'tcx> {
+    pub(crate) fn remove(&mut self, key: OpaqueTypeKey<'tcx>) {
+        match self.opaque_types.remove(&key) {
+            None => bug!("reverted opaque type inference that was never registered"),
+            Some(decl) => assert_ne!(self.opaque_types_vars.remove(decl.concrete_ty), None),
+        }
+    }
+
+    pub fn get_decl(&self, key: &OpaqueTypeKey<'tcx>) -> Option<&OpaqueTypeDecl<'tcx>> {
+        self.opaque_types.get(key)
+    }
+
+    pub fn get_opaque_type_for_infer_var(&self, key: Ty<'tcx>) -> Option<Ty<'tcx>> {
+        self.opaque_types_vars.get(key).copied()
+    }
+
+    pub fn opaque_types(&self) -> OpaqueTypeMap<'tcx> {
+        self.opaque_types.clone()
+    }
+
+    pub fn take_opaque_types(&mut self) -> OpaqueTypeMap<'tcx> {
+        std::mem::take(&mut self.opaque_types)
+    }
+
+    #[inline]
+    pub(crate) fn with_log<'a>(
+        &'a mut self,
+        undo_log: &'a mut InferCtxtUndoLogs<'tcx>,
+    ) -> OpaqueTypeTable<'a, 'tcx> {
+        OpaqueTypeTable { storage: self, undo_log }
+    }
+}
+pub struct OpaqueTypeTable<'a, 'tcx> {
+    storage: &'a mut OpaqueTypeStorage<'tcx>,
+
+    undo_log: &'a mut InferCtxtUndoLogs<'tcx>,
+}
+
+impl<'a, 'tcx> OpaqueTypeTable<'a, 'tcx> {
+    pub fn register(&mut self, key: OpaqueTypeKey<'tcx>, decl: OpaqueTypeDecl<'tcx>) {
+        self.undo_log.push(key);
+        self.storage.opaque_types.insert(key, decl);
+        self.storage.opaque_types_vars.insert(decl.concrete_ty, decl.opaque_type);
+    }
+}
diff --git a/compiler/rustc_infer/src/infer/undo_log.rs b/compiler/rustc_infer/src/infer/undo_log.rs
index 89db8f464b4..3c8a163de54 100644
--- a/compiler/rustc_infer/src/infer/undo_log.rs
+++ b/compiler/rustc_infer/src/infer/undo_log.rs
@@ -4,7 +4,7 @@ use rustc_data_structures::snapshot_vec as sv;
 use rustc_data_structures::undo_log::{Rollback, UndoLogs};
 use rustc_data_structures::unify as ut;
 use rustc_middle::infer::unify_key::RegionVidKey;
-use rustc_middle::ty;
+use rustc_middle::ty::{self, OpaqueTypeKey};
 
 use crate::{
     infer::{region_constraints, type_variable, InferCtxtInner},
@@ -18,6 +18,7 @@ pub struct Snapshot<'tcx> {
 
 /// Records the "undo" data for a single operation that affects some form of inference variable.
 pub(crate) enum UndoLog<'tcx> {
+    OpaqueTypes(OpaqueTypeKey<'tcx>),
     TypeVariables(type_variable::UndoLog<'tcx>),
     ConstUnificationTable(sv::UndoLog<ut::Delegate<ty::ConstVid<'tcx>>>),
     IntUnificationTable(sv::UndoLog<ut::Delegate<ty::IntVid>>),
@@ -42,6 +43,7 @@ macro_rules! impl_from {
 
 // Upcast from a single kind of "undoable action" to the general enum
 impl_from! {
+    OpaqueTypes(OpaqueTypeKey<'tcx>),
     RegionConstraintCollector(region_constraints::UndoLog<'tcx>),
     TypeVariables(type_variable::UndoLog<'tcx>),
 
@@ -64,6 +66,7 @@ impl_from! {
 impl<'tcx> Rollback<UndoLog<'tcx>> for InferCtxtInner<'tcx> {
     fn reverse(&mut self, undo: UndoLog<'tcx>) {
         match undo {
+            UndoLog::OpaqueTypes(key) => self.opaque_type_storage.remove(key),
             UndoLog::TypeVariables(undo) => self.type_variable_storage.reverse(undo),
             UndoLog::ConstUnificationTable(undo) => self.const_unification_storage.reverse(undo),
             UndoLog::IntUnificationTable(undo) => self.int_unification_storage.reverse(undo),
diff --git a/compiler/rustc_typeck/src/check/check.rs b/compiler/rustc_typeck/src/check/check.rs
index 18a0a8767d4..7983b6ea0b2 100644
--- a/compiler/rustc_typeck/src/check/check.rs
+++ b/compiler/rustc_typeck/src/check/check.rs
@@ -647,8 +647,7 @@ fn check_opaque_meets_bounds<'tcx>(
             infcx.instantiate_opaque_types(hir_id, param_env, opaque_ty, span),
         );
 
-        let opaque_type_map = infcx.inner.borrow().opaque_types.clone();
-        for (OpaqueTypeKey { def_id, substs }, opaque_defn) in opaque_type_map {
+        for (OpaqueTypeKey { def_id, substs }, opaque_defn) in infcx.opaque_types() {
             let hidden_type = tcx.type_of(def_id).subst(tcx, substs);
             trace!(?hidden_type);
             match infcx.at(&misc_cause, param_env).eq(opaque_defn.concrete_ty, hidden_type) {
diff --git a/compiler/rustc_typeck/src/check/fallback.rs b/compiler/rustc_typeck/src/check/fallback.rs
index e5da33d113e..d062a0cc55f 100644
--- a/compiler/rustc_typeck/src/check/fallback.rs
+++ b/compiler/rustc_typeck/src/check/fallback.rs
@@ -176,7 +176,7 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
             .type_var_origin(ty)
             .map(|origin| origin.span)
             .unwrap_or(rustc_span::DUMMY_SP);
-        let oty = self.inner.borrow().opaque_types_vars.get(ty).copied();
+        let oty = self.inner.borrow().opaque_type_storage.get_opaque_type_for_infer_var(ty);
         if let Some(opaque_ty) = oty {
             debug!(
                 "fallback_opaque_type_vars(ty={:?}): falling back to opaque type {:?}",
diff --git a/compiler/rustc_typeck/src/check/writeback.rs b/compiler/rustc_typeck/src/check/writeback.rs
index ec88bdf4a37..165ca1574cd 100644
--- a/compiler/rustc_typeck/src/check/writeback.rs
+++ b/compiler/rustc_typeck/src/check/writeback.rs
@@ -498,8 +498,7 @@ impl<'cx, 'tcx> WritebackCx<'cx, 'tcx> {
 
     #[instrument(skip(self, span), level = "debug")]
     fn visit_opaque_types(&mut self, span: Span) {
-        let opaque_types = self.fcx.infcx.inner.borrow().opaque_types.clone();
-        for (opaque_type_key, opaque_defn) in opaque_types {
+        for (opaque_type_key, opaque_defn) in self.fcx.infcx.opaque_types() {
             let hir_id =
                 self.tcx().hir().local_def_id_to_hir_id(opaque_type_key.def_id.expect_local());
             let instantiated_ty = self.resolve(opaque_defn.concrete_ty, &hir_id);