[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.
This commit is contained in:
Andrew Gallant 2025-08-21 14:50:01 -04:00 committed by Andrew Gallant
parent 8ead02e0b1
commit fb2d0af18c
5 changed files with 51 additions and 30 deletions

2
Cargo.lock generated
View file

@ -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",

View file

@ -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 }

View file

@ -8,7 +8,7 @@ pub fn document_symbols_with_options(
file: File,
options: &SymbolsOptions,
) -> Vec<SymbolInfo> {
symbols_for_file(db, file, options)
symbols_for_file(db, file, options).cloned().collect()
}
/// Get all document symbols for a file (hierarchical by default).

View file

@ -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<SymbolInfo> {
) -> impl Iterator<Item = &'db SymbolInfo> {
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<SymbolInfo> {
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<SymbolInfo>,
symbol_stack: Vec<SymbolInfo>,
/// Track if we're currently inside a function (to exclude local variables)
in_function: bool,
/// Options controlling symbol collection
options: &'a SymbolsOptions,
}
impl<'a> SymbolVisitor<'a> {
fn new(options: &'a SymbolsOptions) -> Self {
Self {
symbols: Vec::new(),
symbol_stack: Vec::new(),
in_function: false,
options,
}
options: SymbolsOptionsWithoutQuery,
}
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) => {

View file

@ -28,7 +28,7 @@ pub fn workspace_symbols(db: &dyn Db, query: &str) -> Vec<WorkspaceSymbolInfo> {
for symbol in file_symbols {
results.push(WorkspaceSymbolInfo {
symbol,
symbol: symbol.clone(),
file: *file,
});
}