refactor(workspace): a document_cache module for what a document caches

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,18 +86,17 @@ 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;
 mod sql_function;
 #[cfg(feature = "db")]
 mod statement_execution;
 mod statement_identifier;
-mod tree_sitter;
 
 pub struct WorkspaceServer {
     /// global settings object for this workspace

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<TextRange>,
     diagnostics: Vec<SDiagnostic>,
-    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 {
@@ -220,7 +211,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),
@@ -256,7 +247,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,
@@ -277,7 +268,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)),
@@ -298,14 +289,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)
     }
@@ -323,11 +314,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());
@@ -346,7 +337,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 {
@@ -363,7 +354,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,
                 )
             }),
@@ -377,7 +368,7 @@ impl<'a> StatementMapper<'a> for WithCSTMapper {
     type Output = (DocumentStatement, Arc<tree_sitter::Tree>);
 
     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)
     }
 }
@@ -838,21 +829,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]

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<ParseResult> {
+        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<Result<(), SyntaxDiagnostic>> {
+        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<tree_sitter::Tree> {
+        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);
+    }
+}

crates/pgls_workspace/src/workspace/server/pg_query.rs → crates/pgls_workspace/src/workspace/server/document_cache/parse.rs

@@ -1,14 +1,14 @@
 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},
 };
 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<Arc<pgls_query::NodeEnum>, SyntaxDiagnostic>;

crates/pgls_workspace/src/workspace/server/tree_sitter.rs → 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`].