Use a more lightweight cache for erase_regions

Author: Zoxc

compiler/rustc_middle/src/query/mod.rs

@@ -660,20 +660,6 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    /// Erases regions from `ty` to yield a new type.
-    /// Normally you would just use `tcx.erase_regions(value)`,
-    /// however, which uses this query as a kind of cache.
-    query erase_regions_ty(ty: Ty<'tcx>) -> Ty<'tcx> {
-        // This query is not expected to have input -- as a result, it
-        // is not a good candidates for "replay" because it is essentially a
-        // pure function of its input (and hence the expectation is that
-        // no caller would be green **apart** from just these
-        // queries). Making it anonymous avoids hashing the result, which
-        // may save a bit of time.
-        anon
-        desc { "erasing regions from `{}`", ty }
-    }
-
     query wasm_import_module_map(_: CrateNum) -> &'tcx DefIdMap<String> {
         arena_cache
         desc { "getting wasm import module map" }

compiler/rustc_middle/src/ty/context.rs

@@ -1376,6 +1376,9 @@ pub struct GlobalCtxt<'tcx> {
     /// Common consts, pre-interned for your convenience.
     pub consts: CommonConsts<'tcx>,
 
+    /// A cache for the `erase_regions` function.
+    pub(in crate::ty) erased_region_cache: ShardedHashMap<Ty<'tcx>, Ty<'tcx>>,
+
     /// Hooks to be able to register functions in other crates that can then still
     /// be called from rustc_middle.
     pub(crate) hooks: crate::hooks::Providers,
@@ -1617,6 +1620,7 @@ impl<'tcx> TyCtxt<'tcx> {
             types: common_types,
             lifetimes: common_lifetimes,
             consts: common_consts,
+            erased_region_cache: Default::default(),
             untracked,
             query_system,
             query_kinds,

compiler/rustc_middle/src/ty/erase_regions.rs

@@ -1,21 +1,27 @@
 use tracing::debug;
 
-use crate::query::Providers;
+use crate::dep_graph::DepGraph;
 use crate::ty::{
     self, Ty, TyCtxt, TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
 };
 
-pub(super) fn provide(providers: &mut Providers) {
-    *providers = Providers { erase_regions_ty, ..*providers };
-}
+impl<'tcx> TyCtxt<'tcx> {
+    /// Erases regions from `ty` to yield a new type.
+    pub fn erase_regions_ty(self, ty: Ty<'tcx>) -> Ty<'tcx> {
+        if let Some(ty) = self.erased_region_cache.get(&ty) {
+            return ty;
+        }
 
-fn erase_regions_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Ty<'tcx> {
-    // N.B., use `super_fold_with` here. If we used `fold_with`, it
-    // could invoke the `erase_regions_ty` query recursively.
-    ty.super_fold_with(&mut RegionEraserVisitor { tcx })
-}
+        let result = DepGraph::debug_assert_no_deps(|| {
+            // N.B., use `super_fold_with` here. If we used `fold_with`, it
+            // could invoke the `erase_regions_ty` function recursively.
+            ty.super_fold_with(&mut RegionEraserVisitor { tcx: self })
+        });
+
+        self.erased_region_cache.insert(ty, result);
+        result
+    }
 
-impl<'tcx> TyCtxt<'tcx> {
     /// Returns an equivalent value with all free regions removed (note
     /// that late-bound regions remain, because they are important for
     /// subtyping, but they are anonymized and normalized as well)..

compiler/rustc_middle/src/ty/mod.rs

@@ -2129,7 +2129,6 @@ pub fn ast_uint_ty(uty: UintTy) -> ast::UintTy {
 pub fn provide(providers: &mut Providers) {
     closure::provide(providers);
     context::provide(providers);
-    erase_regions::provide(providers);
     inhabitedness::provide(providers);
     util::provide(providers);
     print::provide(providers);

compiler/rustc_query_system/src/dep_graph/graph.rs

@@ -268,6 +268,14 @@ impl<D: Deps> DepGraph<D> {
         D::with_deps(TaskDepsRef::Forbid, op)
     }
 
+    /// This checks that no dependencies are registered in `op` if debug assertions are enabled.
+    pub fn debug_assert_no_deps<OP, R>(op: OP) -> R
+    where
+        OP: FnOnce() -> R,
+    {
+        if cfg!(debug_assertions) { D::with_deps(TaskDepsRef::Forbid, op) } else { op() }
+    }
+
     #[inline(always)]
     pub fn with_task<Ctxt: HasDepContext<Deps = D>, A: Debug, R>(
         &self,