[ty] Remove excess capacity from more Salsa cached collections (#25411)

Author: MichaReiser

AGENTS.md

@@ -76,3 +76,5 @@ When working on ty, PR titles should start with `[ty]` and be tagged with the `t
 - Run `cargo dev generate-all` after changing configuration options, CLI arguments, lint rules, or environment variable definitions, as these changes require regeneration of schemas, docs, and CLI references.
 - Don't prefix tests with `test_`.
 - Don't separate struct definitions from their `impl` blocks unless the `impl` is deliberately placed in a separate file, as for large structs.
+- The `db` parameter should always be the first, or second, if it's a method taking a `self` argument
+- For Salsa-cached values, avoid retaining excess collection capacity. Prefer boxed slices; otherwise shrink collections that may have spare capacity before returning them. In particular, inspect `HashMap` and `HashSet` values constructed via `extend`, `collect`, explicit reservation, or removal, since those operations can leave capacity that insert-only construction does not.

crates/ty_ide/src/symbols.rs

@@ -792,9 +792,15 @@ impl<'db> SymbolVisitor<'db> {
             remap.push(Some(new_id));
             new.push(symbol);
         }
+
+        new.shrink_to_fit();
+
         FlatSymbols {
             symbols: new,
-            all_names: self.all_origin.map(|_| self.all_names),
+            all_names: self.all_origin.map(|_| {
+                self.all_names.shrink_to_fit();
+                self.all_names
+            }),
         }
     }
 

crates/ty_module_resolver/src/list.rs

@@ -10,7 +10,7 @@ use crate::resolve::{ModuleResolveMode, ResolverContext, resolve_file_module, se
 
 /// List all available modules, including all sub-modules, sorted in lexicographic order.
 pub fn all_modules(db: &dyn Db) -> Vec<Module<'_>> {
-    let mut modules = list_modules(db);
+    let mut modules = list_modules(db).to_vec();
     let mut stack = modules.clone();
     while let Some(module) = stack.pop() {
         for &submodule in module.all_submodules(db) {
@@ -23,8 +23,8 @@ pub fn all_modules(db: &dyn Db) -> Vec<Module<'_>> {
 }
 
 /// List all available top-level modules.
-#[salsa::tracked]
-pub fn list_modules(db: &dyn Db) -> Vec<Module<'_>> {
+#[salsa::tracked(returns(deref))]
+pub fn list_modules(db: &dyn Db) -> Box<[Module<'_>]> {
     let mut modules: BTreeMap<&ModuleName, ListedModule<'_>> = BTreeMap::new();
     for search_path in search_paths(db, ModuleResolveMode::StubsAllowed) {
         for new in list_modules_in(db, SearchPathIngredient::new(db, search_path.clone())) {
@@ -463,7 +463,7 @@ mod tests {
     }
 
     fn sorted_list(db: &dyn Db) -> Vec<Module<'_>> {
-        let mut modules = list_modules(db);
+        let mut modules = list_modules(db).to_vec();
         modules.sort_by(|m1, m2| m1.name(db).cmp(m2.name(db)));
         modules
     }

crates/ty_module_resolver/src/module.rs

@@ -97,9 +97,7 @@ impl<'db> Module<'db> {
     /// The names returned correspond to the "base" name of the module.
     /// That is, `{self.name}.{basename}` should give the full module name.
     pub fn all_submodules(self, db: &'db dyn Db) -> &'db [Module<'db>] {
-        all_submodule_names_for_package(db, self)
-            .as_deref()
-            .unwrap_or_default()
+        all_submodule_names_for_package(db, self).unwrap_or_default()
     }
 }
 
@@ -118,12 +116,11 @@ impl std::fmt::Debug for Module<'_> {
     }
 }
 
-#[allow(clippy::ref_option)]
-#[salsa::tracked(returns(ref), heap_size=ruff_memory_usage::heap_size)]
+#[salsa::tracked(returns(as_deref), heap_size=ruff_memory_usage::heap_size)]
 fn all_submodule_names_for_package<'db>(
     db: &'db dyn Db,
     module: Module<'db>,
-) -> Option<Vec<Module<'db>>> {
+) -> Option<Box<[Module<'db>]>> {
     fn is_submodule(
         is_dir: bool,
         is_file: bool,

crates/ty_module_resolver/src/resolve.rs

@@ -370,7 +370,7 @@ pub fn search_paths(db: &dyn Db, resolve_mode: ModuleResolveMode) -> SearchPathI
 ///
 /// We exclude `__init__.py(i)` dirs to avoid truncating packages.
 #[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
-fn absolute_desperate_search_paths(db: &dyn Db, importing_file: File) -> Option<Vec<SearchPath>> {
+fn absolute_desperate_search_paths(db: &dyn Db, importing_file: File) -> Option<Box<[SearchPath]>> {
     let system = db.system();
     let importing_path = importing_file.path(db).as_system_path()?;
 
@@ -425,7 +425,7 @@ fn absolute_desperate_search_paths(db: &dyn Db, importing_file: File) -> Option<
     if search_paths.is_empty() {
         None
     } else {
-        Some(search_paths)
+        Some(search_paths.into_boxed_slice())
     }
 }
 

crates/ty_python_core/src/ast_ids.rs

@@ -110,9 +110,7 @@ impl AstIdsBuilder {
         self.uses_map.get(&key.into()).copied()
     }
 
-    pub(super) fn finish(mut self) -> AstIds {
-        self.uses_map.shrink_to_fit();
-
+    pub(super) fn finish(self) -> AstIds {
         AstIds {
             uses_map: self.uses_map,
         }

crates/ty_python_core/src/builder.rs

@@ -2230,17 +2230,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         place_tables.shrink_to_fit();
         use_def_maps.shrink_to_fit();
         ast_ids.shrink_to_fit();
-        self.definitions_by_node.shrink_to_fit();
-        self.statements_by_node.shrink_to_fit();
         self.enclosing_lambda_statements.shrink_to_fit();
-        self.collections_by_use.shrink_to_fit();
-        self.uses_by_collection.shrink_to_fit();
-
+        self.imported_modules.shrink_to_fit();
         self.scope_ids_by_scope.shrink_to_fit();
-        self.scopes_by_node.shrink_to_fit();
-        self.generator_functions.shrink_to_fit();
         self.enclosing_snapshots.shrink_to_fit();
 
+        let mut semantic_syntax_errors = self.semantic_syntax_errors.into_inner();
+        semantic_syntax_errors.shrink_to_fit();
+
         SemanticIndex {
             place_tables,
             scopes: self.scopes,
@@ -2258,7 +2255,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             imported_modules: Arc::new(self.imported_modules),
             has_future_annotations: self.has_future_annotations,
             enclosing_snapshots: self.enclosing_snapshots,
-            semantic_syntax_errors: self.semantic_syntax_errors.into_inner(),
+            semantic_syntax_errors,
             generator_functions: self.generator_functions,
             narrowing_alias_predicates: self.alias_predicates,
         }
@@ -3620,10 +3617,12 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
         // context the lambda is being inferred with, and so any statement
         // containing a lambda must be inferable as a standalone statement
         // to avoid large scope-level cycles.
-        for lambda in current_statement.lambda_expressions {
-            self.enclosing_lambda_statements
-                .insert(lambda.into(), standalone_statement);
-        }
+        self.enclosing_lambda_statements.extend(
+            current_statement
+                .lambda_expressions
+                .into_iter()
+                .map(|lambda| (lambda.into(), standalone_statement)),
+        );
 
         // The inferred element type of collection literal depends on uses of
         // the collection in its containing scope, and so each use must be part

crates/ty_python_core/src/use_def.rs

@@ -1683,10 +1683,7 @@ impl<'db> UseDefMapBuilder<'db> {
         self.reachable_symbol_definitions.shrink_to_fit();
         self.reachable_member_definitions.shrink_to_fit();
         self.bindings_by_use.shrink_to_fit();
-        self.multi_bindings_by_use.shrink_to_fit();
         self.range_reachability.shrink_to_fit();
-        self.declarations_by_binding.shrink_to_fit();
-        self.bindings_by_definition.shrink_to_fit();
         self.enclosing_snapshots.shrink_to_fit();
 
         let mut interned_bindings = IndexVec::with_capacity(self.bindings_by_definition.len());
@@ -1808,7 +1805,6 @@ impl<'db> UseDefMapBuilder<'db> {
             interned_ids_by_definition.insert(definition, interned_id);
         }
 
-        interned_ids_by_definition.shrink_to_fit();
         interned_ids_by_definition
     }
 
@@ -1832,7 +1828,6 @@ impl<'db> UseDefMapBuilder<'db> {
             interned_ids_by_binding.insert(binding, interned_id);
         }
 
-        interned_ids_by_binding.shrink_to_fit();
         interned_ids_by_binding
     }
 

crates/ty_python_semantic/src/dunder_all.rs

@@ -1,10 +1,9 @@
-use rustc_hash::FxHashSet;
-
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::statement_visitor::{StatementVisitor, walk_stmt};
 use ruff_python_ast::{self as ast};
+use rustc_hash::FxHashSet;
 use ty_module_resolver::{ModuleName, resolve_module};
 
 use crate::Db;
@@ -197,13 +196,14 @@ impl<'db> DunderAllNamesCollector<'db> {
     ///
     /// Returns [`None`] if `__all__` is not defined in the current module or if it contains
     /// invalid elements.
-    fn into_names(self) -> Option<FxHashSet<Name>> {
+    fn into_names(mut self) -> Option<FxHashSet<Name>> {
         if self.origin.is_none() {
             None
         } else if self.invalid {
             tracing::debug!("Invalid `__all__` in `{}`", self.file.path(self.db));
             None
         } else {
+            self.names.shrink_to_fit();
             Some(self.names)
         }
     }

crates/ty_python_semantic/src/semantic_model.rs

@@ -119,7 +119,8 @@ impl<'db> SemanticModel<'db> {
         let is_typing_extensions_available = self.file.is_stub(self.db)
             || resolve_real_shadowable_module(self.db, self.file, &typing_extensions).is_some();
         list_modules(self.db)
-            .into_iter()
+            .iter()
+            .copied()
             .filter(|module| {
                 is_typing_extensions_available || module.name(self.db) != &typing_extensions
             })