From 52d449a95e606eedf9ff0eb90e35bde638249731 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 26 Oct 2026 15:31:44 +0000 Subject: [PATCH] refactor(workspace): a document_cache module for what a document caches The parse results and the tree-sitter trees of a document's statements move into server/document_cache, behind a DocumentCache that the document owns. No behaviour changes. --- crates/pgls_workspace/src/workspace/server.rs | 7 +-- .../src/workspace/server/document.rs | 55 +++++++---------- .../src/workspace/server/document_cache.rs | 60 +++++++++++++++++++ .../{pg_query.rs => document_cache/parse.rs} | 4 +- .../tree.rs} | 2 +- 5 files changed, 86 insertions(+), 42 deletions(-) create mode 100644 crates/pgls_workspace/src/workspace/server/document_cache.rs rename crates/pgls_workspace/src/workspace/server/{pg_query.rs => document_cache/parse.rs} (98%) rename crates/pgls_workspace/src/workspace/server/{tree_sitter.rs => document_cache/tree.rs} (97%) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 0f60170f6..1a805f96e 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -17,8 +17,6 @@ use document::TypecheckDiagnosticsMapper; use document::{DefaultMapper, Document, ExecuteStatementMapper}; #[cfg(feature = "db")] use futures::{StreamExt, TryStreamExt, stream}; -#[cfg(feature = "db")] -use pg_query::convert_to_positional_params_with_metadata; use pgls_analyse::{ ActionCategory, AnalysisFilter, Applicability, RuleCategories, RuleCategoriesBuilder, RuleFilter, SUPPRESSION_ACTION_CATEGORY, @@ -31,6 +29,8 @@ use pgls_analyser::{ }; use pgls_catalog::resolve::{ColumnSource, Reference, ResolveParams, resolve}; use pgls_catalog::{Catalog, CatalogBase, Origin, Session}; +#[cfg(feature = "db")] +use pgls_lexer::convert_to_positional_params_with_metadata; use std::borrow::Cow; use pgls_catalog::Snapshot; @@ -86,10 +86,10 @@ mod connection_key; #[cfg(feature = "db")] mod connection_manager; pub(crate) mod document; +mod document_cache; mod expand_star; mod formatting; mod migration; -mod pg_query; #[cfg(feature = "db")] mod pgpass; mod snapshot_manager; @@ -97,7 +97,6 @@ mod sql_function; #[cfg(feature = "db")] mod statement_execution; mod statement_identifier; -mod tree_sitter; pub struct WorkspaceServer { /// global settings object for this workspace diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index f44b2c35f..70785aa74 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -7,26 +7,19 @@ use pgls_suppressions::Suppressions; use pgls_text_size::{TextRange, TextSize}; use rustc_hash::FxHashSet; -use super::{ - pg_query::{PgQueryStore, convert_to_positional_params_with_metadata}, - statement_identifier::StatementId, - tree_sitter::TreeSitterStore, -}; +use super::{document_cache::DocumentCache, statement_identifier::StatementId}; pub struct Document { content: String, version: i32, ranges: Vec, diagnostics: Vec, - ast_db: PgQueryStore, - cst_db: TreeSitterStore, + cache: DocumentCache, suppressions: Suppressions, } impl Document { pub fn new(content: String, version: i32) -> Document { - let cst_db = TreeSitterStore::new(); - let ast_db = PgQueryStore::new(); let (ranges, diagnostics) = split_with_diagnostics(&content, None); let suppressions = Suppressions::from(content.as_str()).with_statements(ranges.clone()); @@ -35,8 +28,7 @@ impl Document { diagnostics, content, version, - ast_db, - cst_db, + cache: DocumentCache::new(), suppressions, } } @@ -60,8 +52,7 @@ impl Document { .collect(); let is_present = |id: &StatementId| statements.contains(id.parent_content().unwrap_or(id.content())); - self.ast_db.retain(is_present); - self.cst_db.retain(is_present); + self.cache.retain_statements(is_present); } pub fn suppressions(&self) -> &Suppressions { @@ -219,7 +210,7 @@ where // The statements of a SQL function body follow the function. They are pushed in // reverse, so that they are popped in order. - let embedded = self.parser.ast_db.get_or_cache_embedded(&root_id); + let embedded = self.parser.cache.embedded(&root_id); self.pending_embedded .extend(embedded.iter().rev().map(|embedded| DocumentStatement { id: root_id.create_child(&embedded.sql), @@ -255,7 +246,7 @@ impl<'a> StatementMapper<'a> for ExecuteStatementMapper { ); fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output { - let ast_result = doc.ast_db.get_or_cache_ast(&stmt.id); + let ast_result = doc.cache.parse(&stmt.id); let ast_option = match &*ast_result { Ok(node) => Some(node.clone()), Err(_) => None, @@ -276,7 +267,7 @@ impl<'a> StatementMapper<'a> for FormatStatementMapper { ); fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output { - let ast_result = doc.ast_db.get_or_cache_ast(&stmt.id); + let ast_result = doc.cache.parse(&stmt.id); let result = match &*ast_result { Ok(node) => Ok(node.clone()), Err(diag) => Err(diag.clone().span(stmt.range)), @@ -297,14 +288,14 @@ impl<'a> StatementMapper<'a> for TypecheckDiagnosticsMapper { ); fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output { - let ast_result = doc.ast_db.get_or_cache_ast(&stmt.id); + let ast_result = doc.cache.parse(&stmt.id); let ast_option = match &*ast_result { Ok(node) => Some(node.clone()), Err(_) => None, }; - let cst_result = doc.cst_db.get_or_cache_tree(&stmt.id); + let cst_result = doc.cache.tree(&stmt.id); (stmt.id, stmt.range, ast_option, cst_result) } @@ -322,11 +313,11 @@ impl<'a> StatementMapper<'a> for AnalyserDiagnosticsMapper { fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output { let DocumentStatement { id, range, offsets } = stmt; - let maybe_node = doc.ast_db.get_or_cache_ast(&id); + let maybe_node = doc.cache.parse(&id); let (ast_option, diagnostics) = match &*maybe_node { Ok(node) => { - let plpgsql_result = doc.ast_db.get_or_cache_plpgsql_parse(&id); + let plpgsql_result = doc.cache.plpgsql_parse(&id); if let Some(Err(diag)) = plpgsql_result { // offset the pgpsql diagnostic from the parent statement start let span = diag.location().span.map(|sp| sp + range.start()); @@ -345,7 +336,7 @@ impl<'a> StatementMapper<'a> for AnalyserDiagnosticsMapper { // Statements in the body of a SQL function can reference its parameters. let function = id.parent().and_then(|root| { - let parent = doc.ast_db.get_or_cache_ast(&root); + let parent = doc.cache.parse(&root); let Ok(pgls_query::NodeEnum::CreateFunctionStmt(create_fn)) = parent.as_ref().as_deref() else { @@ -362,7 +353,7 @@ impl<'a> StatementMapper<'a> for AnalyserDiagnosticsMapper { } .with_function(function) .with_identifier_parameters( - convert_to_positional_params_with_metadata(id.content()) + pgls_lexer::convert_to_positional_params_with_metadata(id.content()) .has_identifier_parameters, ) }), @@ -376,7 +367,7 @@ impl<'a> StatementMapper<'a> for WithCSTMapper { type Output = (DocumentStatement, Arc); fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output { - let tree = doc.cst_db.get_or_cache_tree(&stmt.id); + let tree = doc.cache.tree(&stmt.id); (stmt, tree) } } @@ -819,21 +810,15 @@ $$ LANGUAGE plpgsql;"; let mut d = Document::new("select 1;\nselect 2;".to_string(), 1); let first = StatementId::new("select 1;"); let second = StatementId::new("select 2;"); - let first_ast = d.ast_db.get_or_cache_ast(&first); - let second_ast = d.ast_db.get_or_cache_ast(&second); - let first_tree = d.cst_db.get_or_cache_tree(&first); + let first_ast = d.cache.parse(&first); + let second_ast = d.cache.parse(&second); + let first_tree = d.cache.tree(&first); d.update_content("select 1;\nselect 3;".to_string(), 2); - assert!(Arc::ptr_eq(&first_ast, &d.ast_db.get_or_cache_ast(&first))); - assert!(Arc::ptr_eq( - &first_tree, - &d.cst_db.get_or_cache_tree(&first) - )); - assert!(!Arc::ptr_eq( - &second_ast, - &d.ast_db.get_or_cache_ast(&second) - )); + assert!(Arc::ptr_eq(&first_ast, &d.cache.parse(&first))); + assert!(Arc::ptr_eq(&first_tree, &d.cache.tree(&first))); + assert!(!Arc::ptr_eq(&second_ast, &d.cache.parse(&second))); } #[test] diff --git a/crates/pgls_workspace/src/workspace/server/document_cache.rs b/crates/pgls_workspace/src/workspace/server/document_cache.rs new file mode 100644 index 000000000..260230395 --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/document_cache.rs @@ -0,0 +1,60 @@ +//! Everything that is cached for a document. +//! +//! The caches of a document are keyed by the text of its statements, so that an edit only +//! drops what was cached for the statements it changed. + +use std::sync::Arc; + +use pgls_query_ext::{diagnostics::SyntaxDiagnostic, function_body::EmbeddedStatement}; + +use super::StatementId; + +mod parse; +mod tree; + +pub(crate) use parse::ParseResult; +use parse::PgQueryStore; +use tree::TreeSitterStore; + +/// The caches of a document. +pub(crate) struct DocumentCache { + parses: PgQueryStore, + trees: TreeSitterStore, +} + +impl DocumentCache { + pub(crate) fn new() -> Self { + Self { + parses: PgQueryStore::new(), + trees: TreeSitterStore::new(), + } + } + + /// The tree of the statement `id`, or why it doesn't parse. + pub(crate) fn parse(&self, id: &StatementId) -> Arc { + self.parses.get_or_cache_ast(id) + } + + /// The syntax error of the PL/pgSQL body of the function that `id` creates, if it creates a + /// PL/pgSQL function. + pub(crate) fn plpgsql_parse(&self, id: &StatementId) -> Option> { + self.parses.get_or_cache_plpgsql_parse(id) + } + + /// The statements embedded in the top-level statement `id`, see + /// [`PgQueryStore::get_or_cache_embedded`]. + pub(crate) fn embedded(&self, id: &StatementId) -> Arc<[EmbeddedStatement]> { + self.parses.get_or_cache_embedded(id) + } + + /// The tree-sitter tree of the statement `id`. + pub(crate) fn tree(&self, id: &StatementId) -> Arc { + self.trees.get_or_cache_tree(id) + } + + /// Drops what is cached for the statements for which `keep` returns `false`. + pub(crate) fn retain_statements(&self, keep: impl Fn(&StatementId) -> bool) { + self.parses.retain(&keep); + self.trees.retain(&keep); + } +} diff --git a/crates/pgls_workspace/src/workspace/server/pg_query.rs b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs similarity index 98% rename from crates/pgls_workspace/src/workspace/server/pg_query.rs rename to crates/pgls_workspace/src/workspace/server/document_cache/parse.rs index 97e2d4eb3..bf30cfe62 100644 --- a/crates/pgls_workspace/src/workspace/server/pg_query.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/parse.rs @@ -1,6 +1,6 @@ use std::sync::{Arc, LazyLock, Mutex}; -pub use pgls_lexer::{convert_to_positional_params, convert_to_positional_params_with_metadata}; +use pgls_lexer::convert_to_positional_params; use pgls_query_ext::{ diagnostics::*, function_body::{EmbeddedStatement, function_body_location, sql_function_body}, @@ -8,7 +8,7 @@ use pgls_query_ext::{ use regex::Regex; use rustc_hash::FxHashMap; -use super::statement_identifier::StatementId; +use crate::workspace::StatementId; /// The tree of a statement, or why it doesn't parse. pub type ParseResult = Result, SyntaxDiagnostic>; diff --git a/crates/pgls_workspace/src/workspace/server/tree_sitter.rs b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs similarity index 97% rename from crates/pgls_workspace/src/workspace/server/tree_sitter.rs rename to crates/pgls_workspace/src/workspace/server/document_cache/tree.rs index ccbb89881..c699d2f67 100644 --- a/crates/pgls_workspace/src/workspace/server/tree_sitter.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache/tree.rs @@ -2,7 +2,7 @@ use std::sync::{Arc, Mutex}; use rustc_hash::FxHashMap; -use super::statement_identifier::StatementId; +use crate::workspace::StatementId; /// The tree-sitter trees of the statements of a document. They are kept until the statement is /// no longer in the document, see [`TreeSitterStore::retain`].