From fb2d0af18cfd68c0bd129f23148994c33526840a Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Thu, 21 Aug 2025 14:50:01 -0400 Subject: [PATCH] [ty] Optimize "workspace symbols" retrieval Basically, this splits the implementation into two pieces: the first piece does the traversal and finds *all* symbols across the workspace. The second piece does filtering based on a user provided query string. Only the first piece is cached by Salsa. This brings warm "workspace symbols" requests down from 500-600ms to 100-200ms. --- Cargo.lock | 2 + crates/ty_ide/Cargo.toml | 2 + crates/ty_ide/src/document_symbols.rs | 12 ++--- crates/ty_ide/src/symbols.rs | 63 ++++++++++++++++---------- crates/ty_ide/src/workspace_symbols.rs | 2 +- 5 files changed, 51 insertions(+), 30 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 573188cfb2..1e53684735 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4245,12 +4245,14 @@ dependencies = [ "itertools 0.14.0", "regex", "ruff_db", + "ruff_memory_usage", "ruff_python_ast", "ruff_python_parser", "ruff_python_trivia", "ruff_source_file", "ruff_text_size", "rustc-hash", + "salsa", "smallvec", "tracing", "ty_project", diff --git a/crates/ty_ide/Cargo.toml b/crates/ty_ide/Cargo.toml index 5f1b245397..3b423ff1cb 100644 --- a/crates/ty_ide/Cargo.toml +++ b/crates/ty_ide/Cargo.toml @@ -13,6 +13,7 @@ license = { workspace = true } [dependencies] bitflags = { workspace = true } ruff_db = { workspace = true } +ruff_memory_usage = { workspace = true } ruff_python_ast = { workspace = true } ruff_python_parser = { workspace = true } ruff_python_trivia = { workspace = true } @@ -24,6 +25,7 @@ ty_project = { workspace = true, features = ["testing"] } itertools = { workspace = true } regex = { workspace = true } rustc-hash = { workspace = true } +salsa = { workspace = true, features = ["compact_str"] } smallvec = { workspace = true } tracing = { workspace = true } diff --git a/crates/ty_ide/src/document_symbols.rs b/crates/ty_ide/src/document_symbols.rs index e9d5096c07..2edb70241b 100644 --- a/crates/ty_ide/src/document_symbols.rs +++ b/crates/ty_ide/src/document_symbols.rs @@ -8,7 +8,7 @@ pub fn document_symbols_with_options( file: File, options: &SymbolsOptions, ) -> Vec { - symbols_for_file(db, file, options) + symbols_for_file(db, file, options).cloned().collect() } /// Get all document symbols for a file (hierarchical by default). @@ -94,13 +94,13 @@ class MyClass: class_var = 100 typed_class_var: str = 'class_typed' annotated_class_var: float - + def __init__(self): self.instance_var = 0 - + def public_method(self): return self.instance_var - + def _private_method(self): pass @@ -256,10 +256,10 @@ def standalone_function(): " class OuterClass: OUTER_CONSTANT = 100 - + def outer_method(self): return self.OUTER_CONSTANT - + class InnerClass: def inner_method(self): pass diff --git a/crates/ty_ide/src/symbols.rs b/crates/ty_ide/src/symbols.rs index ef4614b3fc..4d4e72ea1f 100644 --- a/crates/ty_ide/src/symbols.rs +++ b/crates/ty_ide/src/symbols.rs @@ -125,47 +125,64 @@ impl SymbolKind { } } -pub(crate) fn symbols_for_file( - db: &dyn Db, +pub(crate) fn symbols_for_file<'db>( + db: &'db dyn Db, file: File, options: &SymbolsOptions, -) -> Vec { +) -> impl Iterator { assert!( !options.hierarchical || options.query_string.is_none(), "Cannot use hierarchical mode with a query string" ); + let ingredient = SymbolsOptionsWithoutQuery { + hierarchical: options.hierarchical, + global_only: options.global_only, + }; + symbols_for_file_inner(db, file, ingredient) + .iter() + .filter(|symbol| { + let Some(ref query) = options.query_string else { + return true; + }; + query.is_match(symbol) + }) +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct SymbolsOptionsWithoutQuery { + hierarchical: bool, + global_only: bool, +} + +#[salsa::tracked(returns(deref))] +fn symbols_for_file_inner<'db>( + db: &'db dyn Db, + file: File, + options: SymbolsOptionsWithoutQuery, +) -> Vec { let parsed = parsed_module(db, file); let module = parsed.load(db); - let mut visitor = SymbolVisitor::new(options); + let mut visitor = SymbolVisitor { + symbols: vec![], + symbol_stack: vec![], + in_function: false, + options, + }; visitor.visit_body(&module.syntax().body); - let mut symbols = visitor.symbols; - if let Some(ref query) = options.query_string { - symbols.retain(|symbol| query.is_match(symbol)); - } - symbols + visitor.symbols } -struct SymbolVisitor<'a> { +struct SymbolVisitor { symbols: Vec, symbol_stack: Vec, /// Track if we're currently inside a function (to exclude local variables) in_function: bool, - /// Options controlling symbol collection - options: &'a SymbolsOptions, + options: SymbolsOptionsWithoutQuery, } -impl<'a> SymbolVisitor<'a> { - fn new(options: &'a SymbolsOptions) -> Self { - Self { - symbols: Vec::new(), - symbol_stack: Vec::new(), - in_function: false, - options, - } - } - +impl SymbolVisitor { fn visit_body(&mut self, body: &[Stmt]) { for stmt in body { self.visit_stmt(stmt); @@ -205,7 +222,7 @@ impl<'a> SymbolVisitor<'a> { } } -impl SourceOrderVisitor<'_> for SymbolVisitor<'_> { +impl SourceOrderVisitor<'_> for SymbolVisitor { fn visit_stmt(&mut self, stmt: &Stmt) { match stmt { Stmt::FunctionDef(func_def) => { diff --git a/crates/ty_ide/src/workspace_symbols.rs b/crates/ty_ide/src/workspace_symbols.rs index 5bda7090e5..96a65e92a1 100644 --- a/crates/ty_ide/src/workspace_symbols.rs +++ b/crates/ty_ide/src/workspace_symbols.rs @@ -28,7 +28,7 @@ pub fn workspace_symbols(db: &dyn Db, query: &str) -> Vec { for symbol in file_symbols { results.push(WorkspaceSymbolInfo { - symbol, + symbol: symbol.clone(), file: *file, }); }