Auto merge of #17639 - Veykril:salsa-perf, r=Veykril

Some more small salsa memory improvements

This does limit our lru limits to 2^16 but if you want to set them higher than that you might as well not set them at all. Also makes `LRU` opt-in per query now, allowing us to drop all the unnecessary LRU stuff for most queries

Author: bors

crates/base-db/src/lib.rs

@@ -42,9 +42,9 @@ pub trait Upcast<T: ?Sized> {
     fn upcast(&self) -> &T;
 }
 
-pub const DEFAULT_FILE_TEXT_LRU_CAP: usize = 16;
-pub const DEFAULT_PARSE_LRU_CAP: usize = 128;
-pub const DEFAULT_BORROWCK_LRU_CAP: usize = 2024;
+pub const DEFAULT_FILE_TEXT_LRU_CAP: u16 = 16;
+pub const DEFAULT_PARSE_LRU_CAP: u16 = 128;
+pub const DEFAULT_BORROWCK_LRU_CAP: u16 = 2024;
 
 pub trait FileLoader {
     /// Text of the file.
@@ -59,6 +59,7 @@ pub trait FileLoader {
 #[salsa::query_group(SourceDatabaseStorage)]
 pub trait SourceDatabase: FileLoader + std::fmt::Debug {
     /// Parses the file into the syntax tree.
+    #[salsa::lru]
     fn parse(&self, file_id: EditionedFileId) -> Parse<ast::SourceFile>;
 
     /// Returns the set of errors obtained from parsing the file including validation errors.
@@ -105,6 +106,7 @@ pub trait SourceDatabaseExt: SourceDatabase {
     #[salsa::input]
     fn compressed_file_text(&self, file_id: FileId) -> Arc<[u8]>;
 
+    #[salsa::lru]
     fn file_text(&self, file_id: FileId) -> Arc<str>;
 
     /// Path to a file, relative to the root of its source root.

crates/hir-expand/src/db.rs

@@ -65,7 +65,7 @@ pub trait ExpandDatabase: SourceDatabase {
     #[salsa::transparent]
     fn parse_or_expand_with_err(&self, file_id: HirFileId) -> ExpandResult<Parse<SyntaxNode>>;
     /// Implementation for the macro case.
-    // This query is LRU cached
+    #[salsa::lru]
     fn parse_macro_expansion(
         &self,
         macro_file: MacroFileId,

crates/hir-ty/src/db.rs

@@ -61,6 +61,7 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     ) -> Result<Arc<MirBody>, MirLowerError>;
 
     #[salsa::invoke(crate::mir::borrowck_query)]
+    #[salsa::lru]
     fn borrowck(&self, def: DefWithBodyId) -> Result<Arc<[BorrowckResult]>, MirLowerError>;
 
     #[salsa::invoke(crate::consteval::const_eval_query)]

crates/ide-db/src/lib.rs

@@ -145,7 +145,7 @@ impl Default for RootDatabase {
 }
 
 impl RootDatabase {
-    pub fn new(lru_capacity: Option<usize>) -> RootDatabase {
+    pub fn new(lru_capacity: Option<u16>) -> RootDatabase {
         let mut db = RootDatabase { storage: ManuallyDrop::new(salsa::Storage::default()) };
         db.set_crate_graph_with_durability(Default::default(), Durability::HIGH);
         db.set_proc_macros_with_durability(Default::default(), Durability::HIGH);
@@ -161,7 +161,7 @@ impl RootDatabase {
         self.set_expand_proc_attr_macros_with_durability(true, Durability::HIGH);
     }
 
-    pub fn update_base_query_lru_capacities(&mut self, lru_capacity: Option<usize>) {
+    pub fn update_base_query_lru_capacities(&mut self, lru_capacity: Option<u16>) {
         let lru_capacity = lru_capacity.unwrap_or(base_db::DEFAULT_PARSE_LRU_CAP);
         base_db::FileTextQuery.in_db_mut(self).set_lru_capacity(DEFAULT_FILE_TEXT_LRU_CAP);
         base_db::ParseQuery.in_db_mut(self).set_lru_capacity(lru_capacity);
@@ -170,7 +170,7 @@ impl RootDatabase {
         hir::db::BorrowckQuery.in_db_mut(self).set_lru_capacity(base_db::DEFAULT_BORROWCK_LRU_CAP);
     }
 
-    pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, usize>) {
+    pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, u16>) {
         use hir::db as hir_db;
 
         base_db::FileTextQuery.in_db_mut(self).set_lru_capacity(DEFAULT_FILE_TEXT_LRU_CAP);
@@ -192,135 +192,6 @@ impl RootDatabase {
                 .copied()
                 .unwrap_or(base_db::DEFAULT_BORROWCK_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::MacroArgQuery
-            hir_db::DeclMacroExpanderQuery
-            // hir_db::MacroExpandQuery
-            hir_db::ExpandProcMacroQuery
-            hir_db::ParseMacroExpansionErrorQuery
-
-            // DefDatabase
-            hir_db::FileItemTreeQuery
-            hir_db::BlockDefMapQuery
-            hir_db::StructDataWithDiagnosticsQuery
-            hir_db::UnionDataWithDiagnosticsQuery
-            hir_db::EnumDataQuery
-            hir_db::EnumVariantDataWithDiagnosticsQuery
-            hir_db::ImplDataWithDiagnosticsQuery
-            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::FieldsAttrsQuery
-            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::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::InternCoroutineQuery
-            hir_db::AssociatedTyDataQuery
-            hir_db::TraitDatumQuery
-            hir_db::AdtDatumQuery
-            hir_db::ImplDatumQuery
-            hir_db::FnDefDatumQuery
-            hir_db::FnDefVarianceQuery
-            hir_db::AdtVarianceQuery
-            hir_db::AssociatedTyValueQuery
-            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
-        ];
     }
 }
 

crates/ide/src/lib.rs

@@ -161,19 +161,19 @@ pub struct AnalysisHost {
 }
 
 impl AnalysisHost {
-    pub fn new(lru_capacity: Option<usize>) -> AnalysisHost {
+    pub fn new(lru_capacity: Option<u16>) -> AnalysisHost {
         AnalysisHost { db: RootDatabase::new(lru_capacity) }
     }
 
     pub fn with_database(db: RootDatabase) -> AnalysisHost {
         AnalysisHost { db }
     }
 
-    pub fn update_lru_capacity(&mut self, lru_capacity: Option<usize>) {
+    pub fn update_lru_capacity(&mut self, lru_capacity: Option<u16>) {
         self.db.update_base_query_lru_capacities(lru_capacity);
     }
 
-    pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, usize>) {
+    pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, u16>) {
         self.db.update_lru_capacities(lru_capacities);
     }
 

crates/load-cargo/src/lib.rs

@@ -377,7 +377,7 @@ fn load_crate_graph(
 ) -> RootDatabase {
     let ProjectWorkspace { toolchain, target_layout, .. } = ws;
 
-    let lru_cap = std::env::var("RA_LRU_CAP").ok().and_then(|it| it.parse::<usize>().ok());
+    let lru_cap = std::env::var("RA_LRU_CAP").ok().and_then(|it| it.parse::<u16>().ok());
     let mut db = RootDatabase::new(lru_cap);
     let mut analysis_change = ChangeWithProcMacros::new();
 

crates/rust-analyzer/src/config.rs

@@ -283,9 +283,9 @@ config_data! {
         linkedProjects: Vec<ManifestOrProjectJson> = vec![],
 
         /// Number of syntax trees rust-analyzer keeps in memory. Defaults to 128.
-        lru_capacity: Option<usize>                 = None,
+        lru_capacity: Option<u16>                 = None,
         /// Sets the LRU capacity of the specified queries.
-        lru_query_capacities: FxHashMap<Box<str>, usize> = FxHashMap::default(),
+        lru_query_capacities: FxHashMap<Box<str>, u16> = FxHashMap::default(),
 
         /// These proc-macros will be ignored when trying to expand them.
         ///
@@ -1743,11 +1743,11 @@ impl Config {
         extra_env
     }
 
-    pub fn lru_parse_query_capacity(&self) -> Option<usize> {
+    pub fn lru_parse_query_capacity(&self) -> Option<u16> {
         self.lru_capacity().to_owned()
     }
 
-    pub fn lru_query_capacities_config(&self) -> Option<&FxHashMap<Box<str>, usize>> {
+    pub fn lru_query_capacities_config(&self) -> Option<&FxHashMap<Box<str>, u16>> {
         self.lru_query_capacities().is_empty().not().then(|| self.lru_query_capacities())
     }
 
@@ -2926,7 +2926,7 @@ fn field_props(field: &str, ty: &str, doc: &[&str], default: &str) -> serde_json
         "FxHashMap<String, String>" => set! {
             "type": "object",
         },
-        "FxHashMap<Box<str>, usize>" => set! {
+        "FxHashMap<Box<str>, u16>" => set! {
             "type": "object",
         },
         "FxHashMap<String, Option<String>>" => set! {
@@ -2936,6 +2936,11 @@ fn field_props(field: &str, ty: &str, doc: &[&str], default: &str) -> serde_json
             "type": ["null", "integer"],
             "minimum": 0,
         },
+        "Option<u16>" => set! {
+            "type": ["null", "integer"],
+            "minimum": 0,
+            "maximum": 65535,
+        },
         "Option<String>" => set! {
             "type": ["null", "string"],
         },

crates/salsa/salsa-macros/src/query_group.rs

@@ -53,7 +53,11 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
                         num_storages += 1;
                     }
                     "dependencies" => {
-                        storage = QueryStorage::Dependencies;
+                        storage = QueryStorage::LruDependencies;
+                        num_storages += 1;
+                    }
+                    "lru" => {
+                        storage = QueryStorage::LruMemoized;
                         num_storages += 1;
                     }
                     "input" => {
@@ -235,7 +239,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
 
         queries_with_storage.push(fn_name);
 
-        let tracing = if let QueryStorage::Memoized = query.storage {
+        let tracing = if let QueryStorage::Memoized | QueryStorage::LruMemoized = query.storage {
             let s = format!("{trait_name}::{fn_name}");
             Some(quote! {
                 let _p = tracing::debug_span!(#s, #(#key_names = tracing::field::debug(&#key_names)),*).entered();
@@ -376,8 +380,9 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
 
         let storage = match &query.storage {
             QueryStorage::Memoized => quote!(salsa::plumbing::MemoizedStorage<Self>),
-            QueryStorage::Dependencies => {
-                quote!(salsa::plumbing::DependencyStorage<Self>)
+            QueryStorage::LruMemoized => quote!(salsa::plumbing::LruMemoizedStorage<Self>),
+            QueryStorage::LruDependencies => {
+                quote!(salsa::plumbing::LruDependencyStorage<Self>)
             }
             QueryStorage::Input if query.keys.is_empty() => {
                 quote!(salsa::plumbing::UnitInputStorage<Self>)
@@ -724,7 +729,8 @@ impl Query {
 #[derive(Debug, Clone, PartialEq, Eq)]
 enum QueryStorage {
     Memoized,
-    Dependencies,
+    LruDependencies,
+    LruMemoized,
     Input,
     Interned,
     InternedLookup { intern_query_type: Ident },
@@ -739,7 +745,9 @@ impl QueryStorage {
             | QueryStorage::Interned
             | QueryStorage::InternedLookup { .. }
             | QueryStorage::Transparent => false,
-            QueryStorage::Memoized | QueryStorage::Dependencies => true,
+            QueryStorage::Memoized | QueryStorage::LruMemoized | QueryStorage::LruDependencies => {
+                true
+            }
         }
     }
 }