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),