From e0b3ca42f230a36b4b2e695c92a35e4a51c0fca1 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 24 Oct 2026 21:11:18 +0000 Subject: [PATCH] perf(analyser): share statement trees with the workspace instead of copying them The workspace copied the parse tree of every statement into the analyser (AnalysableStatement owned its tree), and the analyser copied them all once more for the file context. Analysing a file of 5,000 statements peaked at twice its cached trees. AnalysableStatement::root is now an Arc of the tree that the workspace caches, and the file context borrows them. The other statement mappers share the trees too; the formatter copies a tree only when it replaces a function body in it. --- crates/pgls_analyser/src/lib.rs | 11 ++++++----- crates/pgls_analyser/src/linter_context.rs | 8 ++++---- crates/pgls_workspace/src/workspace/server.rs | 8 ++++++-- .../src/workspace/server/document.rs | 13 +++++++++---- .../src/workspace/server/pg_query.rs | 15 ++++++++------- crates/pgls_workspace/tests/memory_bounds.rs | 4 +++- 6 files changed, 36 insertions(+), 23 deletions(-) diff --git a/crates/pgls_analyser/src/lib.rs b/crates/pgls_analyser/src/lib.rs index 479374548..985581b64 100644 --- a/crates/pgls_analyser/src/lib.rs +++ b/crates/pgls_analyser/src/lib.rs @@ -77,7 +77,9 @@ pub fn replacement_of_removed_rule(name: &str) -> Option<&'static str> { #[derive(Debug)] pub struct AnalysableStatement { - pub root: pgls_query::NodeEnum, + /// Shared with the workspace's cache of parsed statements, so that analysing a file doesn't + /// copy every statement's tree. + pub root: Arc, pub range: TextRange, /// The statement text, used for precise spans of typecheck findings. pub sql: Option, @@ -90,9 +92,9 @@ pub struct AnalysableStatement { } impl AnalysableStatement { - pub fn new(root: pgls_query::NodeEnum, range: TextRange) -> Self { + pub fn new(root: impl Into>, range: TextRange) -> Self { Self { - root, + root: root.into(), range, sql: None, function: None, @@ -191,8 +193,7 @@ impl<'a> Analyser<'a> { .catalog_base .as_ref() .map(|base| base.snapshot().as_ref()); - let roots: Vec = - params.stmts.iter().map(|s| s.root.clone()).collect(); + let roots: Vec<&pgls_query::NodeEnum> = params.stmts.iter().map(|s| &*s.root).collect(); let rules: Vec<_> = self .registry diff --git a/crates/pgls_analyser/src/linter_context.rs b/crates/pgls_analyser/src/linter_context.rs index ab61173b5..6115f90a1 100644 --- a/crates/pgls_analyser/src/linter_context.rs +++ b/crates/pgls_analyser/src/linter_context.rs @@ -180,14 +180,14 @@ impl<'a> StatementContext<'a> { /// /// In a file with migration sections, the file context is that of the current section. pub struct AnalysedFileContext<'a> { - pub stmts: &'a [pgls_query::NodeEnum], + pub stmts: &'a [&'a pgls_query::NodeEnum], pos: usize, catalog: Catalog, session: Session, } impl<'a> AnalysedFileContext<'a> { - pub fn new(stmts: &'a [pgls_query::NodeEnum], catalog: Catalog, session: Session) -> Self { + pub fn new(stmts: &'a [&'a pgls_query::NodeEnum], catalog: Catalog, session: Session) -> Self { Self { stmts, pos: 0, @@ -196,7 +196,7 @@ impl<'a> AnalysedFileContext<'a> { } } - pub fn previous_stmts(&self) -> &[pgls_query::NodeEnum] { + pub fn previous_stmts(&self) -> &[&'a pgls_query::NodeEnum] { &self.stmts[0..self.pos] } @@ -238,7 +238,7 @@ impl<'a> AnalysedFileContext<'a> { /// in the bodies of SQL functions are only analysed, not executed. pub fn next(&mut self, applies_effects: bool) { if applies_effects && self.pos < self.stmts.len() { - let stmt = &self.stmts[self.pos]; + let stmt = self.stmts[self.pos]; self.catalog.apply(stmt, self.session.search_path()); self.session.apply(stmt); } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 6879b89d1..b5910362a 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -1728,7 +1728,11 @@ impl Workspace for WorkspaceServer { Ok(ast) => { let mut ast = ast; if let Some(formatted_sql_fn_body) = formatted_sql_fn_bodies.get(&id) { - sql_function::set_formatted_sql_fn_body(&mut ast, formatted_sql_fn_body); + // Copies the tree only for the functions whose body is replaced. + sql_function::set_formatted_sql_fn_body( + Arc::make_mut(&mut ast), + formatted_sql_fn_body, + ); } match pgls_pretty_print::format_statement(&ast, &text, &config) { @@ -1952,7 +1956,7 @@ impl Workspace for WorkspaceServer { let markdown_blocks = pgls_hover::on_hover(pgls_hover::OnHoverParams { ts_tree: &ts_tree, snapshot: &snapshot, - ast: maybe_ast.as_ref(), + ast: maybe_ast.as_deref(), position: position_in_stmt, stmt_sql: stmt_id.content(), references: &references, diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index df86fb4db..c2521e872 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -258,7 +258,12 @@ impl<'a> StatementMapper<'a> for DefaultMapper { pub struct ExecuteStatementMapper; impl<'a> StatementMapper<'a> for ExecuteStatementMapper { - type Output = (StatementId, TextRange, String, Option); + type Output = ( + StatementId, + TextRange, + String, + Option>, + ); fn map(&self, parser: &'a Document, id: StatementId, range: TextRange) -> Self::Output { let ast_result = parser.ast_db.get_or_cache_ast(&id); @@ -277,7 +282,7 @@ impl<'a> StatementMapper<'a> for FormatStatementMapper { StatementId, TextRange, String, - Result, + Result, pgls_query_ext::diagnostics::SyntaxDiagnostic>, ); fn map(&self, parser: &'a Document, id: StatementId, range: TextRange) -> Self::Output { @@ -296,7 +301,7 @@ impl<'a> StatementMapper<'a> for TypecheckDiagnosticsMapper { type Output = ( StatementId, TextRange, - Option, + Option>, Arc, Option, ); @@ -409,7 +414,7 @@ impl<'a> StatementMapper<'a> for WithCSTandASTMapper { StatementId, TextRange, Arc, - Option, + Option>, ); fn map(&self, parser: &'a Document, id: StatementId, range: TextRange) -> Self::Output { diff --git a/crates/pgls_workspace/src/workspace/server/pg_query.rs b/crates/pgls_workspace/src/workspace/server/pg_query.rs index cb50a891c..e3e5b3e38 100644 --- a/crates/pgls_workspace/src/workspace/server/pg_query.rs +++ b/crates/pgls_workspace/src/workspace/server/pg_query.rs @@ -8,10 +8,14 @@ use rustc_hash::FxHashMap; use super::statement_identifier::StatementId; +/// The tree of a statement, or why it doesn't parse. +pub type ParseResult = Result, SyntaxDiagnostic>; + /// The parse results of the statements of a document. They are kept until the statement is /// no longer in the document, see [`PgQueryStore::retain`]. pub struct PgQueryStore { - ast_db: Mutex>>>, + /// The tree is in an `Arc` of its own, so that the analyser can share it. + ast_db: Mutex>>, plpgsql_db: Mutex>>, } @@ -40,10 +44,7 @@ impl PgQueryStore { self.plpgsql_db.lock().unwrap().retain(|id, _| keep(id)); } - pub fn get_or_cache_ast( - &self, - statement: &StatementId, - ) -> Arc> { + pub fn get_or_cache_ast(&self, statement: &StatementId) -> Arc { let mut cache = self.ast_db.lock().unwrap(); if let Some(existing) = cache.get(statement) { @@ -54,7 +55,7 @@ impl PgQueryStore { pgls_query::parse(&convert_to_positional_params(statement.content())) .map_err(SyntaxDiagnostic::from) .and_then(|ast| { - ast.into_root().ok_or_else(|| { + ast.into_root().map(Arc::new).ok_or_else(|| { SyntaxDiagnostic::new("No root node found in parse result", None) }) }), @@ -69,7 +70,7 @@ impl PgQueryStore { ) -> Option> { let ast = self.get_or_cache_ast(statement); - let create_fn = match ast.as_ref() { + let create_fn = match ast.as_ref().as_deref() { Ok(pgls_query::NodeEnum::CreateFunctionStmt(node)) => node, _ => return None, }; diff --git a/crates/pgls_workspace/tests/memory_bounds.rs b/crates/pgls_workspace/tests/memory_bounds.rs index 596657b33..38571cda6 100644 --- a/crates/pgls_workspace/tests/memory_bounds.rs +++ b/crates/pgls_workspace/tests/memory_bounds.rs @@ -170,6 +170,7 @@ async fn open_documents_stay_within_the_budget(test_db: PgPool) { // The CLI: each file is opened, checked and closed. let reset = reset_peak(); + let cli_started = std::time::Instant::now(); for (document, content) in contents.iter().enumerate() { let path = PgLSPath::new(format!("cli_{document}.sql")); workspace @@ -190,13 +191,14 @@ async fn open_documents_stay_within_the_budget(test_db: PgPool) { .unwrap(); workspace.close_file(CloseFileParams { path }).unwrap(); } + let cli_duration = cli_started.elapsed(); let cli_peak = peak().saturating_sub(reset); let cli_retained = current().saturating_sub(baseline); println!( "{documents} documents, {:.1} MiB of SQL: one document peak {:.1} MiB, retained {:.1} MiB; \ all open: peak {:.1} MiB, retained {:.1} MiB; retained after closing {:.1} MiB; \ - CLI peak {:.1} MiB, retained {:.1} MiB", + CLI peak {:.1} MiB, retained {:.1} MiB, {cli_duration:.1?}", mib(text), mib(first.0), mib(first.1),