Auto merge of #14410 - Veykril:query-lru-capacities, r=Veykril

internal: Add config to specifiy lru capacities for all queries

Might help figuring out what queries should be limited by LRU by default, as currently we only limit `parse`, `parse_macro_expansion` and `macro_expand`.

8 files changed, 202 insertions, 9 deletions

diff --git a/crates/hir-def/src/attr.rs b/crates/hir-def/src/attr.rs
index 200072c172e..e4d3636ce5e 100644
--- a/crates/hir-def/src/attr.rs
+++ b/crates/hir-def/src/attr.rs
@@ -88,6 +88,7 @@ impl Attrs {
         db: &dyn DefDatabase,
         e: EnumId,
     ) -> Arc<ArenaMap<LocalEnumVariantId, Attrs>> {
+        let _p = profile::span("variants_attrs_query");
         // FIXME: There should be some proper form of mapping between item tree enum variant ids and hir enum variant ids
         let mut res = ArenaMap::default();
 
@@ -114,6 +115,7 @@ impl Attrs {
         db: &dyn DefDatabase,
         v: VariantId,
     ) -> Arc<ArenaMap<LocalFieldId, Attrs>> {
+        let _p = profile::span("fields_attrs_query");
         // FIXME: There should be some proper form of mapping between item tree field ids and hir field ids
         let mut res = ArenaMap::default();
 
@@ -253,6 +255,7 @@ impl Attrs {
 
 impl AttrsWithOwner {
     pub(crate) fn attrs_query(db: &dyn DefDatabase, def: AttrDefId) -> Self {
+        let _p = profile::span("attrs_query");
         // FIXME: this should use `Trace` to avoid duplication in `source_map` below
         let raw_attrs = match def {
             AttrDefId::ModuleId(module) => {
diff --git a/crates/ide-db/src/lib.rs b/crates/ide-db/src/lib.rs
index f9b8a502d99..64f22261d21 100644
--- a/crates/ide-db/src/lib.rs
+++ b/crates/ide-db/src/lib.rs
@@ -141,16 +141,175 @@ impl RootDatabase {
         db.set_local_roots_with_durability(Default::default(), Durability::HIGH);
         db.set_library_roots_with_durability(Default::default(), Durability::HIGH);
         db.set_enable_proc_attr_macros(false);
-        db.update_lru_capacity(lru_capacity);
+        db.update_parse_query_lru_capacity(lru_capacity);
         db
     }
 
-    pub fn update_lru_capacity(&mut self, lru_capacity: Option<usize>) {
+    pub fn update_parse_query_lru_capacity(&mut self, lru_capacity: Option<usize>) {
         let lru_capacity = lru_capacity.unwrap_or(base_db::DEFAULT_LRU_CAP);
         base_db::ParseQuery.in_db_mut(self).set_lru_capacity(lru_capacity);
         hir::db::ParseMacroExpansionQuery.in_db_mut(self).set_lru_capacity(lru_capacity);
         hir::db::MacroExpandQuery.in_db_mut(self).set_lru_capacity(lru_capacity);
     }
+
+    pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, usize>) {
+        use hir::db as hir_db;
+
+        base_db::ParseQuery.in_db_mut(self).set_lru_capacity(
+            lru_capacities.get(stringify!(ParseQuery)).copied().unwrap_or(base_db::DEFAULT_LRU_CAP),
+        );
+        hir_db::ParseMacroExpansionQuery.in_db_mut(self).set_lru_capacity(
+            lru_capacities
+                .get(stringify!(ParseMacroExpansionQuery))
+                .copied()
+                .unwrap_or(base_db::DEFAULT_LRU_CAP),
+        );
+        hir_db::MacroExpandQuery.in_db_mut(self).set_lru_capacity(
+            lru_capacities
+                .get(stringify!(MacroExpandQuery))
+                .copied()
+                .unwrap_or(base_db::DEFAULT_LRU_CAP),
+        );
+
+        macro_rules! update_lru_capacity_per_query {
+            ($( $module:ident :: $query:ident )*) => {$(
+                if let Some(&cap) = lru_capacities.get(stringify!($query)) {
+                    $module::$query.in_db_mut(self).set_lru_capacity(cap);
+                }
+            )*}
+        }
+        update_lru_capacity_per_query![
+            // SourceDatabase
+            // base_db::ParseQuery
+            // base_db::CrateGraphQuery
+            // base_db::ProcMacrosQuery
+
+            // SourceDatabaseExt
+            // base_db::FileTextQuery
+            // base_db::FileSourceRootQuery
+            // base_db::SourceRootQuery
+            base_db::SourceRootCratesQuery
+
+            // ExpandDatabase
+            hir_db::AstIdMapQuery
+            // hir_db::ParseMacroExpansionQuery
+            // hir_db::InternMacroCallQuery
+            hir_db::MacroArgTextQuery
+            hir_db::MacroDefQuery
+            // hir_db::MacroExpandQuery
+            hir_db::ExpandProcMacroQuery
+            hir_db::MacroExpandErrorQuery
+            hir_db::HygieneFrameQuery
+
+            // DefDatabase
+            hir_db::FileItemTreeQuery
+            hir_db::CrateDefMapQueryQuery
+            hir_db::BlockDefMapQuery
+            hir_db::StructDataQuery
+            hir_db::StructDataWithDiagnosticsQuery
+            hir_db::UnionDataQuery
+            hir_db::UnionDataWithDiagnosticsQuery
+            hir_db::EnumDataQuery
+            hir_db::EnumDataWithDiagnosticsQuery
+            hir_db::ImplDataQuery
+            hir_db::ImplDataWithDiagnosticsQuery
+            hir_db::TraitDataQuery
+            hir_db::TraitDataWithDiagnosticsQuery
+            hir_db::TraitAliasDataQuery
+            hir_db::TypeAliasDataQuery
+            hir_db::FunctionDataQuery
+            hir_db::ConstDataQuery
+            hir_db::StaticDataQuery
+            hir_db::Macro2DataQuery
+            hir_db::MacroRulesDataQuery
+            hir_db::ProcMacroDataQuery
+            hir_db::BodyWithSourceMapQuery
+            hir_db::BodyQuery
+            hir_db::ExprScopesQuery
+            hir_db::GenericParamsQuery
+            hir_db::VariantsAttrsQuery
+            hir_db::FieldsAttrsQuery
+            hir_db::VariantsAttrsSourceMapQuery
+            hir_db::FieldsAttrsSourceMapQuery
+            hir_db::AttrsQuery
+            hir_db::CrateLangItemsQuery
+            hir_db::LangItemQuery
+            hir_db::ImportMapQuery
+            hir_db::FieldVisibilitiesQuery
+            hir_db::FunctionVisibilityQuery
+            hir_db::ConstVisibilityQuery
+            hir_db::CrateSupportsNoStdQuery
+
+            // HirDatabase
+            hir_db::InferQueryQuery
+            hir_db::MirBodyQuery
+            hir_db::BorrowckQuery
+            hir_db::TyQuery
+            hir_db::ValueTyQuery
+            hir_db::ImplSelfTyQuery
+            hir_db::ConstParamTyQuery
+            hir_db::ConstEvalQuery
+            hir_db::ConstEvalDiscriminantQuery
+            hir_db::ImplTraitQuery
+            hir_db::FieldTypesQuery
+            hir_db::LayoutOfAdtQuery
+            hir_db::TargetDataLayoutQuery
+            hir_db::CallableItemSignatureQuery
+            hir_db::ReturnTypeImplTraitsQuery
+            hir_db::GenericPredicatesForParamQuery
+            hir_db::GenericPredicatesQuery
+            hir_db::TraitEnvironmentQuery
+            hir_db::GenericDefaultsQuery
+            hir_db::InherentImplsInCrateQuery
+            hir_db::InherentImplsInBlockQuery
+            hir_db::IncoherentInherentImplCratesQuery
+            hir_db::TraitImplsInCrateQuery
+            hir_db::TraitImplsInBlockQuery
+            hir_db::TraitImplsInDepsQuery
+            // hir_db::InternCallableDefQuery
+            // hir_db::InternLifetimeParamIdQuery
+            // hir_db::InternImplTraitIdQuery
+            // hir_db::InternTypeOrConstParamIdQuery
+            // hir_db::InternClosureQuery
+            // hir_db::InternGeneratorQuery
+            hir_db::AssociatedTyDataQuery
+            hir_db::TraitDatumQuery
+            hir_db::StructDatumQuery
+            hir_db::ImplDatumQuery
+            hir_db::FnDefDatumQuery
+            hir_db::FnDefVarianceQuery
+            hir_db::AdtVarianceQuery
+            hir_db::AssociatedTyValueQuery
+            hir_db::TraitSolveQueryQuery
+            hir_db::ProgramClausesForChalkEnvQuery
+
+            // SymbolsDatabase
+            symbol_index::ModuleSymbolsQuery
+            symbol_index::LibrarySymbolsQuery
+            // symbol_index::LocalRootsQuery
+            // symbol_index::LibraryRootsQuery
+
+            // LineIndexDatabase
+            crate::LineIndexQuery
+
+            // InternDatabase
+            // hir_db::InternFunctionQuery
+            // hir_db::InternStructQuery
+            // hir_db::InternUnionQuery
+            // hir_db::InternEnumQuery
+            // hir_db::InternConstQuery
+            // hir_db::InternStaticQuery
+            // hir_db::InternTraitQuery
+            // hir_db::InternTraitAliasQuery
+            // hir_db::InternTypeAliasQuery
+            // hir_db::InternImplQuery
+            // hir_db::InternExternBlockQuery
+            // hir_db::InternBlockQuery
+            // hir_db::InternMacro2Query
+            // hir_db::InternProcMacroQuery
+            // hir_db::InternMacroRulesQuery
+        ];
+    }
 }
 
 impl salsa::ParallelDatabase for RootDatabase {
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 8477a8e6228..f70496451d2 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -67,7 +67,7 @@ use ide_db::{
         salsa::{self, ParallelDatabase},
         CrateOrigin, Env, FileLoader, FileSet, SourceDatabase, VfsPath,
     },
-    symbol_index, LineIndexDatabase,
+    symbol_index, FxHashMap, LineIndexDatabase,
 };
 use syntax::SourceFile;
 
@@ -154,7 +154,11 @@ impl AnalysisHost {
     }
 
     pub fn update_lru_capacity(&mut self, lru_capacity: Option<usize>) {
-        self.db.update_lru_capacity(lru_capacity);
+        self.db.update_parse_query_lru_capacity(lru_capacity);
+    }
+
+    pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, usize>) {
+        self.db.update_lru_capacities(lru_capacities);
     }
 
     /// Returns a snapshot of the current state, which you can query for
diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index e917f107f3b..06fb95515ae 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/crates/rust-analyzer/src/config.rs
@@ -7,7 +7,7 @@
 //! configure the server itself, feature flags are passed into analysis, and
 //! tweak things like automatic insertion of `()` in completions.
 
-use std::{fmt, iter, path::PathBuf};
+use std::{fmt, iter, ops::Not, path::PathBuf};
 
 use flycheck::FlycheckConfig;
 use ide::{
@@ -418,6 +418,8 @@ config_data! {
 
         /// Number of syntax trees rust-analyzer keeps in memory. Defaults to 128.
         lru_capacity: Option<usize>                 = "null",
+        /// Sets the LRU capacity of the specified queries.
+        lru_query_capacities: FxHashMap<Box<str>, usize> = "{}",
 
         /// Whether to show `can't find Cargo.toml` error message.
         notifications_cargoTomlNotFound: bool      = "true",
@@ -1085,10 +1087,14 @@ impl Config {
         extra_env
     }
 
-    pub fn lru_capacity(&self) -> Option<usize> {
+    pub fn lru_parse_query_capacity(&self) -> Option<usize> {
         self.data.lru_capacity
     }
 
+    pub fn lru_query_capacities(&self) -> Option<&FxHashMap<Box<str>, usize>> {
+        self.data.lru_query_capacities.is_empty().not().then(|| &self.data.lru_query_capacities)
+    }
+
     pub fn proc_macro_srv(&self) -> Option<(AbsPathBuf, /* is path explicitly set */ bool)> {
         if !self.data.procMacro_enable {
             return None;
@@ -2024,6 +2030,9 @@ fn field_props(field: &str, ty: &str, doc: &[&str], default: &str) -> serde_json
         "FxHashMap<String, String>" => set! {
             "type": "object",
         },
+        "FxHashMap<Box<str>, usize>" => set! {
+            "type": "object",
+        },
         "Option<usize>" => set! {
             "type": ["null", "integer"],
             "minimum": 0,
diff --git a/crates/rust-analyzer/src/global_state.rs b/crates/rust-analyzer/src/global_state.rs
index 82588c458c5..616d8fca5f3 100644
--- a/crates/rust-analyzer/src/global_state.rs
+++ b/crates/rust-analyzer/src/global_state.rs
@@ -141,7 +141,10 @@ impl GlobalState {
             Handle { handle, receiver }
         };
 
-        let analysis_host = AnalysisHost::new(config.lru_capacity());
+        let mut analysis_host = AnalysisHost::new(config.lru_parse_query_capacity());
+        if let Some(capacities) = config.lru_query_capacities() {
+            analysis_host.update_lru_capacities(capacities);
+        }
         let (flycheck_sender, flycheck_receiver) = unbounded();
         let mut this = GlobalState {
             sender,
diff --git a/crates/rust-analyzer/src/reload.rs b/crates/rust-analyzer/src/reload.rs
index c6bd3f0d9cc..b111813a7cd 100644
--- a/crates/rust-analyzer/src/reload.rs
+++ b/crates/rust-analyzer/src/reload.rs
@@ -74,8 +74,13 @@ impl GlobalState {
     pub(crate) fn update_configuration(&mut self, config: Config) {
         let _p = profile::span("GlobalState::update_configuration");
         let old_config = mem::replace(&mut self.config, Arc::new(config));
-        if self.config.lru_capacity() != old_config.lru_capacity() {
-            self.analysis_host.update_lru_capacity(self.config.lru_capacity());
+        if self.config.lru_parse_query_capacity() != old_config.lru_parse_query_capacity() {
+            self.analysis_host.update_lru_capacity(self.config.lru_parse_query_capacity());
+        }
+        if self.config.lru_query_capacities() != old_config.lru_query_capacities() {
+            self.analysis_host.update_lru_capacities(
+                &self.config.lru_query_capacities().cloned().unwrap_or_default(),
+            );
         }
         if self.config.linked_projects() != old_config.linked_projects() {
             self.fetch_workspaces_queue.request_op("linked projects changed".to_string(), ())
diff --git a/docs/user/generated_config.adoc b/docs/user/generated_config.adoc
index 6937a7ed9a2..12dfe394f1d 100644
--- a/docs/user/generated_config.adoc
+++ b/docs/user/generated_config.adoc
@@ -639,6 +639,11 @@ Elements must be paths pointing to `Cargo.toml`,
 --
 Number of syntax trees rust-analyzer keeps in memory. Defaults to 128.
 --
+[[rust-analyzer.lru.query.capacities]]rust-analyzer.lru.query.capacities (default: `{}`)::
++
+--
+Sets the LRU capacity of the specified queries.
+--
 [[rust-analyzer.notifications.cargoTomlNotFound]]rust-analyzer.notifications.cargoTomlNotFound (default: `true`)::
 +
 --
diff --git a/editors/code/package.json b/editors/code/package.json
index 0332fe30251..81fa97269a9 100644
--- a/editors/code/package.json
+++ b/editors/code/package.json
@@ -1252,6 +1252,11 @@
                     ],
                     "minimum": 0
                 },
+                "rust-analyzer.lru.query.capacities": {
+                    "markdownDescription": "Sets the LRU capacity of the specified queries.",
+                    "default": {},
+                    "type": "object"
+                },
                 "rust-analyzer.notifications.cargoTomlNotFound": {
                     "markdownDescription": "Whether to show `can't find Cargo.toml` error message.",
                     "default": true,