Cache file_to_def in SourceToDefCtx

Author: Veykril

crates/hir/src/semantics.rs

@@ -129,7 +129,7 @@ pub struct Semantics<'db, DB> {
 
 pub struct SemanticsImpl<'db> {
     pub db: &'db dyn HirDatabase,
-    s2d_cache: RefCell<(SourceToDefCache, FxHashMap<MacroFileId, hir_expand::ExpansionInfo>)>,
+    s2d_cache: RefCell<SourceToDefCache>,
     /// Rootnode to HirFileId cache
     root_to_file_cache: RefCell<FxHashMap<SyntaxNode, HirFileId>>,
     /// MacroCall to its expansion's MacroFileId cache
@@ -719,7 +719,8 @@ impl<'db> SemanticsImpl<'db> {
             let macro_file = invoc.as_macro_file();
             let expansion_info = {
                 self.with_ctx(|ctx| {
-                    ctx.expansion_info_cache
+                    ctx.cache
+                        .expansion_info_cache
                         .entry(macro_file)
                         .or_insert_with(|| {
                             let exp_info = macro_file.expansion_info(self.db.upcast());
@@ -806,7 +807,8 @@ impl<'db> SemanticsImpl<'db> {
         let process_expansion_for_token = |stack: &mut Vec<_>, macro_file| {
             let InMacroFile { file_id, value: mapped_tokens } = self.with_ctx(|ctx| {
                 Some(
-                    ctx.expansion_info_cache
+                    ctx.cache
+                        .expansion_info_cache
                         .entry(macro_file)
                         .or_insert_with(|| {
                             let exp_info = macro_file.expansion_info(self.db.upcast());
@@ -1086,6 +1088,7 @@ impl<'db> SemanticsImpl<'db> {
 
                     self.with_ctx(|ctx| {
                         let expansion_info = ctx
+                            .cache
                             .expansion_info_cache
                             .entry(macro_file)
                             .or_insert_with(|| macro_file.expansion_info(self.db.upcast()));
@@ -1364,8 +1367,7 @@ impl<'db> SemanticsImpl<'db> {
     }
 
     fn with_ctx<F: FnOnce(&mut SourceToDefCtx<'_, '_>) -> T, T>(&self, f: F) -> T {
-        let (dynmap_cache, expansion_info_cache) = &mut *self.s2d_cache.borrow_mut();
-        let mut ctx = SourceToDefCtx { db: self.db, dynmap_cache, expansion_info_cache };
+        let mut ctx = SourceToDefCtx { db: self.db, cache: &mut self.s2d_cache.borrow_mut() };
         f(&mut ctx)
     }
 
@@ -1375,7 +1377,7 @@ impl<'db> SemanticsImpl<'db> {
     }
 
     fn file_to_module_defs(&self, file: FileId) -> impl Iterator<Item = Module> {
-        self.with_ctx(|ctx| ctx.file_to_def(file)).into_iter().map(Module::from)
+        self.with_ctx(|ctx| ctx.file_to_def(file).to_owned()).into_iter().map(Module::from)
     }
 
     pub fn scope(&self, node: &SyntaxNode) -> Option<SemanticsScope<'db>> {
@@ -1653,6 +1655,7 @@ fn macro_call_to_macro_id(
                 }
                 HirFileIdRepr::MacroFile(macro_file) => {
                     let expansion_info = ctx
+                        .cache
                         .expansion_info_cache
                         .entry(macro_file)
                         .or_insert_with(|| macro_file.expansion_info(ctx.db.upcast()));
@@ -1668,6 +1671,7 @@ fn macro_call_to_macro_id(
                 }
                 HirFileIdRepr::MacroFile(macro_file) => {
                     let expansion_info = ctx
+                        .cache
                         .expansion_info_cache
                         .entry(macro_file)
                         .or_insert_with(|| macro_file.expansion_info(ctx.db.upcast()));

crates/hir/src/semantics/source_to_def.rs

@@ -112,31 +112,37 @@ use syntax::{
 
 use crate::{db::HirDatabase, InFile};
 
-pub(super) type SourceToDefCache = FxHashMap<(ChildContainer, HirFileId), DynMap>;
+#[derive(Default)]
+pub(super) struct SourceToDefCache {
+    pub(super) dynmap_cache: FxHashMap<(ChildContainer, HirFileId), DynMap>,
+    pub(super) expansion_info_cache: FxHashMap<MacroFileId, ExpansionInfo>,
+    pub(super) file_to_def_cache: FxHashMap<FileId, SmallVec<[ModuleId; 1]>>,
+}
 
-pub(super) struct SourceToDefCtx<'a, 'dyn_cache> {
-    pub(super) db: &'a dyn HirDatabase,
-    pub(super) dynmap_cache: &'dyn_cache mut SourceToDefCache,
-    pub(super) expansion_info_cache: &'a mut FxHashMap<MacroFileId, ExpansionInfo>,
+pub(super) struct SourceToDefCtx<'db, 'cache> {
+    pub(super) db: &'db dyn HirDatabase,
+    pub(super) cache: &'cache mut SourceToDefCache,
 }
 
 impl SourceToDefCtx<'_, '_> {
-    pub(super) fn file_to_def(&self, file: FileId) -> SmallVec<[ModuleId; 1]> {
+    pub(super) fn file_to_def(&mut self, file: FileId) -> &SmallVec<[ModuleId; 1]> {
         let _p = tracing::span!(tracing::Level::INFO, "SourceToDefCtx::file_to_def").entered();
-        let mut mods = SmallVec::new();
-        for &crate_id in self.db.relevant_crates(file).iter() {
-            // Note: `mod` declarations in block modules cannot be supported here
-            let crate_def_map = self.db.crate_def_map(crate_id);
-            mods.extend(
-                crate_def_map
-                    .modules_for_file(file)
-                    .map(|local_id| crate_def_map.module_id(local_id)),
-            )
-        }
-        if mods.is_empty() {
-            // FIXME: detached file
-        }
-        mods
+        self.cache.file_to_def_cache.entry(file).or_insert_with(|| {
+            let mut mods = SmallVec::new();
+            for &crate_id in self.db.relevant_crates(file).iter() {
+                // Note: `mod` declarations in block modules cannot be supported here
+                let crate_def_map = self.db.crate_def_map(crate_id);
+                mods.extend(
+                    crate_def_map
+                        .modules_for_file(file)
+                        .map(|local_id| crate_def_map.module_id(local_id)),
+                )
+            }
+            if mods.is_empty() {
+                // FIXME: detached file
+            }
+            mods
+        })
     }
 
     pub(super) fn module_to_def(&mut self, src: InFile<&ast::Module>) -> Option<ModuleId> {
@@ -166,7 +172,7 @@ impl SourceToDefCtx<'_, '_> {
         Some(def_map.module_id(child_id))
     }
 
-    pub(super) fn source_file_to_def(&self, src: InFile<&ast::SourceFile>) -> Option<ModuleId> {
+    pub(super) fn source_file_to_def(&mut self, src: InFile<&ast::SourceFile>) -> Option<ModuleId> {
         let _p = tracing::span!(tracing::Level::INFO, "source_file_to_def").entered();
         let file_id = src.file_id.original_file(self.db.upcast());
         self.file_to_def(file_id).first().copied()
@@ -325,7 +331,8 @@ impl SourceToDefCtx<'_, '_> {
 
     fn cache_for(&mut self, container: ChildContainer, file_id: HirFileId) -> &DynMap {
         let db = self.db;
-        self.dynmap_cache
+        self.cache
+            .dynmap_cache
             .entry((container, file_id))
             .or_insert_with(|| container.child_by_source(db, file_id))
     }
@@ -421,6 +428,7 @@ impl SourceToDefCtx<'_, '_> {
                 let macro_file = node.file_id.macro_file()?;
 
                 let expansion_info = this
+                    .cache
                     .expansion_info_cache
                     .entry(macro_file)
                     .or_insert_with(|| macro_file.expansion_info(this.db.upcast()));