perf(analyser): share statement trees with the workspace instead of copying them

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<pgls_query::NodeEnum>,
     pub range: TextRange,
     /// The statement text, used for precise spans of typecheck findings.
     pub sql: Option<String>,
@@ -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<Arc<pgls_query::NodeEnum>>, 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<pgls_query::NodeEnum> =
-            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

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);
         }

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,

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<pgls_query::NodeEnum>);
+    type Output = (
+        StatementId,
+        TextRange,
+        String,
+        Option<Arc<pgls_query::NodeEnum>>,
+    );
 
     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<pgls_query::NodeEnum, pgls_query_ext::diagnostics::SyntaxDiagnostic>,
+        Result<Arc<pgls_query::NodeEnum>, 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<pgls_query::NodeEnum>,
+        Option<Arc<pgls_query::NodeEnum>>,
         Arc<tree_sitter::Tree>,
         Option<SQLFunctionSignature>,
     );
@@ -409,7 +414,7 @@ impl<'a> StatementMapper<'a> for WithCSTandASTMapper {
         StatementId,
         TextRange,
         Arc<tree_sitter::Tree>,
-        Option<pgls_query::NodeEnum>,
+        Option<Arc<pgls_query::NodeEnum>>,
     );
 
     fn map(&self, parser: &'a Document, id: StatementId, range: TextRange) -> Self::Output {

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<Arc<pgls_query::NodeEnum>, 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<FxHashMap<StatementId, Arc<Result<pgls_query::NodeEnum, SyntaxDiagnostic>>>>,
+    /// The tree is in an `Arc` of its own, so that the analyser can share it.
+    ast_db: Mutex<FxHashMap<StatementId, Arc<ParseResult>>>,
     plpgsql_db: Mutex<FxHashMap<StatementId, Result<(), SyntaxDiagnostic>>>,
 }
 
@@ -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<Result<pgls_query::NodeEnum, SyntaxDiagnostic>> {
+    pub fn get_or_cache_ast(&self, statement: &StatementId) -> Arc<ParseResult> {
         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<Result<(), SyntaxDiagnostic>> {
         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,
         };

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