diff --git a/crates/pgls_analyser/src/lib.rs b/crates/pgls_analyser/src/lib.rs index 123635b1e..4c61253fc 100644 --- a/crates/pgls_analyser/src/lib.rs +++ b/crates/pgls_analyser/src/lib.rs @@ -27,7 +27,7 @@ pub use linter_registry::{ pub use linter_rule::{FixEdit, LinterDiagnostic, LinterRule, RuleFix}; pub use migration::{Direction, MigrationSection}; use walk::is_plain_query; -pub use walk::{WalkParams, walk_statements}; +pub use walk::{Walk, WalkParams, WalkState, walk_statements}; // For convenience in macros and rule files - keep these shorter names pub use LinterDiagnostic as RuleDiagnostic; @@ -243,14 +243,25 @@ impl<'a> Analyser<'a> { }) .collect(); + let mut statement_ranges: Vec = params + .stmts + .iter() + .filter(|stmt| stmt.function.is_none()) + .map(|stmt| stmt.range) + .chain(params.unparsable.iter().map(|skipped| skipped.range)) + .collect(); + statement_ranges.sort_unstable_by_key(|range| range.start()); let walk = WalkParams { stmts: ¶ms.stmts, unparsable: ¶ms.unparsable, + statement_ranges: &statement_ranges, catalog_base: params.catalog_base.clone(), search_path: ¶ms.search_path, sections: ¶ms.sections, end: TextSize::from(u32::MAX), best_effort: false, + from: None, + checkpoints: &[], }; walk_statements(walk, |stmt, file_context| { let root = &*stmt.root; diff --git a/crates/pgls_analyser/src/migration.rs b/crates/pgls_analyser/src/migration.rs index 6156e9447..8d909ef6e 100644 --- a/crates/pgls_analyser/src/migration.rs +++ b/crates/pgls_analyser/src/migration.rs @@ -2,7 +2,7 @@ //! file and run each of them on its own. use pgls_catalog::TransactionMode; -use pgls_text_size::TextRange; +use pgls_text_size::{TextRange, TextSize}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Direction { @@ -33,4 +33,21 @@ impl MigrationSection { TransactionMode::Autocommit } } + + /// The queries that run in an implicit transaction block: outside a transaction, those with + /// more than one statement. `statement_ranges` are the ranges of the top-level statements of + /// the file, parsed or not, in order. + /// + /// Cost: O(log statements) per query. + pub fn implicit_blocks<'a>( + &'a self, + statement_ranges: &'a [TextRange], + ) -> impl Iterator + 'a { + let starting_before = |offset: TextSize| { + statement_ranges.partition_point(|statement| statement.start() < offset) + }; + self.queries.iter().copied().filter(move |query| { + !self.transaction && starting_before(query.end()) - starting_before(query.start()) > 1 + }) + } } diff --git a/crates/pgls_analyser/src/walk.rs b/crates/pgls_analyser/src/walk.rs index cf2993ba6..e70fda8be 100644 --- a/crates/pgls_analyser/src/walk.rs +++ b/crates/pgls_analyser/src/walk.rs @@ -17,6 +17,10 @@ pub struct WalkParams<'a> { pub stmts: &'a [AnalysableStatement], /// The top-level statements that don't parse, in file order. pub unparsable: &'a [UnparsableStatement], + /// The ranges of all top-level statements of the file, parsed or not, in order. The walk + /// counts the statements of a section and of its queries with them, also when it is only + /// given some of the statements, see [`MigrationSection::implicit_blocks`]. + pub statement_ranges: &'a [TextRange], /// The database snapshot the catalog starts from. pub catalog_base: Option>, /// Explicit search path at the start of the file. @@ -30,12 +34,67 @@ pub struct WalkParams<'a> { /// like a `DO` block, see [`Catalog::best_effort`]. Editor features show what names most /// likely refer to; lint only reports what is certain. pub best_effort: bool, + /// Resumes a walk where it stopped instead of starting at the beginning of the file. The + /// statements are then those from [`WalkState::at`] on. + pub from: Option, + /// The starts of top-level statements where to record the state of the walk, in file + /// order, see [`Walk::checkpoints`]. + pub checkpoints: &'a [TextSize], +} + +/// The state of a walk before a top-level statement: the catalog and the session after the +/// statements before it. It only depends on the text before the statement, the section the +/// statement is in, and the parameters of the walk other than the statements. +#[derive(Clone)] +pub struct WalkState { + at: TextSize, + /// The start of the migration section that `at` is in, or zero before the first section. + section_start: TextSize, + catalog: Catalog, + session: Session, + /// In a down section, the catalog after the up sections before it. + up_catalog: Option, + /// The start of the query whose implicit transaction block the session is in. + implicit_block: Option, +} + +impl WalkState { + /// The start of the statement before which the walk stopped. + pub fn at(&self) -> TextSize { + self.at + } + + /// Whether `sections` put the statement at [`Self::at`] in the section it was in when the + /// walk stopped. A marker inserted right before the statement changes that, even though the + /// text before the statement's old start is the same. + pub fn fits(&self, sections: &[MigrationSection]) -> bool { + section_start(sections, self.at) == self.section_start + } +} + +/// The start of the section that `offset` is in, or zero before the first section. +fn section_start(sections: &[MigrationSection], offset: TextSize) -> TextSize { + sections + .iter() + .take_while(|section| section.range.start() <= offset) + .last() + .map_or(TextSize::from(0), |section| section.range.start()) +} + +pub struct Walk { + /// The catalog as a statement at [`WalkParams::end`] sees it. + pub catalog: Catalog, + /// The state of the walk at each of [`WalkParams::checkpoints`] that it got to: before the + /// first statement that starts at or after it. Resuming from it walks like walking from the + /// start of the file, as long as the text before it doesn't change. + pub checkpoints: Vec, } /// Walks the statements like the database runs them: the catalog and the session follow the /// statements, so that each statement sees the objects that the statements before it create, /// change or drop. `visit` gets each statement with the file context as it is before it. -/// Returns the catalog as a statement at [`WalkParams::end`] sees it. +/// Returns the catalog as a statement at [`WalkParams::end`] sees it, and the states at the +/// checkpoints. /// /// - Each migration section runs in a session and transaction of its own. A down section runs /// after the up sections before it, and its changes don't carry into the next section. @@ -45,13 +104,13 @@ pub struct WalkParams<'a> { /// Such a statement is visited and applied to the session (its locks and its effect on the /// transaction), but if it may change the catalog, it taints it instead of being applied to it. /// -/// The file context only knows the statements that the walk is given. When they stop before the -/// end of the file, [`AnalysedFileContext::stmt_count`] and the implicit transaction blocks of -/// the last query are those of the statements given, which only rules rely on. +/// The file context only knows the statements that the walk is given, except for +/// [`AnalysedFileContext::stmt_count`] and the queries that run in an implicit transaction +/// block, which [`WalkParams::statement_ranges`] tell. pub fn walk_statements( params: WalkParams<'_>, mut visit: impl FnMut(&AnalysableStatement, &AnalysedFileContext), -) -> Catalog { +) -> Walk { let roots: Vec<&pgls_query::NodeEnum> = params.stmts.iter().map(|s| &*s.root).collect(); // The catalog after the up sections so far. @@ -60,80 +119,105 @@ pub fn walk_statements( catalog = catalog.best_effort(); } let mut unparsable = params.unparsable.iter().peekable(); + let mut from = params.from; + // The walk can't record the states before where it resumes. + let resumes_at = from.as_ref().map_or(TextSize::from(0), WalkState::at); + let mut checkpoint_at = params + .checkpoints + .iter() + .copied() + .filter(|at| *at >= resumes_at) + .peekable(); + let mut checkpoints = Vec::new(); for unit in migration_units(params.stmts, params.sections) { if unit.start >= params.end { break; } + if from.as_ref().is_some_and(|from| unit.end <= from.at) { + continue; + } let is_down = unit .section .is_some_and(|section| section.direction == Direction::Down); - let unit_catalog = if is_down { - catalog.clone() - } else { - std::mem::replace(&mut catalog, Catalog::new(None)) + let mut implicit_block = None; + let (unit_catalog, session) = match from.take() { + Some(from) => { + debug_assert_eq!( + from.section_start, unit.start, + "the state fits the sections" + ); + if let Some(up_catalog) = from.up_catalog { + catalog = up_catalog; + } + implicit_block = from.implicit_block; + (from.catalog, from.session) + } + None => { + let unit_catalog = if is_down { + catalog.clone() + } else { + std::mem::replace(&mut catalog, Catalog::new(None)) + }; + // Each section runs in a session and transaction of its own. + let session = Session::new(params.search_path.to_vec()).with_transaction_mode( + unit.section + .map_or(TransactionMode::Unknown, MigrationSection::transaction_mode), + ); + (unit_catalog, session) + } }; - // Each section runs in a session and transaction of its own. - let session = Session::new(params.search_path.to_vec()).with_transaction_mode( - unit.section - .map_or(TransactionMode::Unknown, MigrationSection::transaction_mode), - ); let stmts = ¶ms.stmts[unit.stmts.clone()]; - let unit_unparsable = { - let starts = |end: TextSize| { + let last_end = { + let unparsable_before = |end: TextSize| { params .unparsable .partition_point(|skipped| skipped.range.start() < end) }; - ¶ms.unparsable[starts(unit.start)..starts(unit.end)] + params.unparsable[unparsable_before(unit.start)..unparsable_before(unit.end)] + .iter() + .map(|skipped| skipped.range.end()) + .chain(stmts.last().map(|stmt| stmt.range.end())) + .max() + .unwrap_or_default() }; - let last_end = unit_unparsable - .iter() - .map(|skipped| skipped.range.end()) - .chain(stmts.last().map(|stmt| stmt.range.end())) - .max() - .unwrap_or_default(); - // The starts of the statements the unit runs, in order, including those that don't - // parse, but not those in function bodies. - let mut starts: Vec = stmts - .iter() - .filter(|stmt| stmt.function.is_none()) - .map(|stmt| stmt.range.start()) - .chain(unit_unparsable.iter().map(|skipped| skipped.range.start())) - .collect(); - starts.sort_unstable(); - let stmt_count = |range: TextRange| { - starts.partition_point(|start| *start < range.end()) - - starts.partition_point(|start| *start < range.start()) + let statements_before = |offset: TextSize| { + params + .statement_ranges + .partition_point(|statement| statement.start() < offset) }; - // Outside a transaction, the queries with more than one statement run in an implicit - // transaction block. - let implicit_blocks: Vec = unit - .section - .filter(|section| !section.transaction) - .map(|section| { - section - .queries - .iter() - .copied() - .filter(|query| stmt_count(*query) > 1) - .collect() - }) - .unwrap_or_default(); // The blocks that don't end before the statement, which comes after those before it. - let mut implicit_blocks = implicit_blocks.into_iter().peekable(); + let mut implicit_blocks = unit + .section + .into_iter() + .flat_map(|section| section.implicit_blocks(params.statement_ranges)) + .peekable(); let mut file_context = AnalysedFileContext::new(&roots[unit.stmts.clone()], unit_catalog, session) .with_positions( stmts.iter().map(|stmt| stmt.range.start()).collect(), last_end, ) - .with_stmt_count(stmt_count(TextRange::new(unit.start, unit.end))); - let mut implicit_block = None; + .with_stmt_count(statements_before(unit.end) - statements_before(unit.start)); + let state = |at: TextSize, + file_context: &AnalysedFileContext, + catalog: &Catalog, + implicit_block: Option| WalkState { + at, + section_start: unit.start, + catalog: file_context.catalog().clone(), + session: file_context.session().clone(), + up_catalog: is_down.then(|| catalog.clone()), + implicit_block, + }; for stmt in stmts { if stmt.range.start() >= params.end { break; } + while let Some(at) = checkpoint_at.next_if(|at| stmt.range.start() >= *at) { + skip_unparsable(&mut unparsable, at, &mut file_context); + checkpoints.push(state(at, &file_context, &catalog, implicit_block)); + } let runs = stmt.function.is_none(); if runs { while implicit_blocks @@ -142,8 +226,8 @@ pub fn walk_statements( {} let block = implicit_blocks .peek() - .copied() - .filter(|block| block.contains(stmt.range.start())); + .filter(|block| block.contains(stmt.range.start())) + .map(|block| block.start()); if block != implicit_block { if implicit_block.is_some() { file_context.end_implicit_block(); @@ -168,16 +252,26 @@ pub fn walk_statements( file_context.next(runs, !placeholders); } + while let Some(at) = checkpoint_at.next_if(|at| *at < unit.end && *at <= params.end) { + skip_unparsable(&mut unparsable, at, &mut file_context); + checkpoints.push(state(at, &file_context, &catalog, implicit_block)); + } skip_unparsable(&mut unparsable, unit.end.min(params.end), &mut file_context); // The section that contains `end` decides what a statement there sees. if params.end < unit.end { - return file_context.into_catalog(); + return Walk { + catalog: file_context.into_catalog(), + checkpoints, + }; } if !is_down { catalog = file_context.into_catalog(); } } - catalog + Walk { + catalog, + checkpoints, + } } /// Moves past the statements that don't parse and start before `before`. Those that may have @@ -318,7 +412,8 @@ mod tests { fn walk_queries(sql: &str, queries: Vec) -> (Vec<(String, bool)>, Duration) { let mut stmts = Vec::new(); let mut unparsable = Vec::new(); - for range in pgls_statement_splitter::split(sql).ranges { + let statement_ranges = pgls_statement_splitter::split(sql).ranges; + for &range in &statement_ranges { let text = &sql[range]; match pgls_query::parse(text) { Ok(parsed) => stmts.push( @@ -342,11 +437,14 @@ mod tests { WalkParams { stmts: &stmts, unparsable: &unparsable, + statement_ranges: &statement_ranges, catalog_base: None, search_path: &[], sections: §ions, end: TextSize::of(sql), best_effort: false, + from: None, + checkpoints: &[], }, |stmt, file_context| { visited.push(( diff --git a/crates/pgls_catalog/src/catalog/ddl/alter_table_stmt.rs b/crates/pgls_catalog/src/catalog/ddl/alter_table_stmt.rs index 7da318ab3..f8e72ee55 100644 --- a/crates/pgls_catalog/src/catalog/ddl/alter_table_stmt.rs +++ b/crates/pgls_catalog/src/catalog/ddl/alter_table_stmt.rs @@ -44,7 +44,7 @@ pub(super) fn apply_alter_table_stmt(c: &mut Catalog, n: &AlterTableStmt, search } type_info.origin = Origin::File; let key = key(&type_info.schema, &type_info.name); - c.types.insert(key, Entry::Defined(type_info)); + c.types.insert(key, Entry::defined(type_info)); } return; } diff --git a/crates/pgls_catalog/src/catalog/ddl/rename_stmt.rs b/crates/pgls_catalog/src/catalog/ddl/rename_stmt.rs index a1cb5f033..4a1e94c0d 100644 --- a/crates/pgls_catalog/src/catalog/ddl/rename_stmt.rs +++ b/crates/pgls_catalog/src/catalog/ddl/rename_stmt.rs @@ -40,7 +40,7 @@ pub(super) fn apply_rename_stmt(c: &mut Catalog, n: &RenameStmt, search_path: &[ rename_in(&mut type_info.attributes, &n.subname, &n.newname); type_info.origin = Origin::File; let key = key(&type_info.schema, &type_info.name); - c.types.insert(key, Entry::Defined(type_info)); + c.types.insert(key, Entry::defined(type_info)); } } } diff --git a/crates/pgls_catalog/src/catalog/materialize.rs b/crates/pgls_catalog/src/catalog/materialize.rs index 8bfa034ac..6687441d8 100644 --- a/crates/pgls_catalog/src/catalog/materialize.rs +++ b/crates/pgls_catalog/src/catalog/materialize.rs @@ -78,7 +78,7 @@ impl Catalog { for (key, entry) in &self.functions { match entry { Entry::Defined(overloads) => { - for function in overloads { + for function in overloads.iter() { if function.origin == Origin::File { builder.add_function(key, function); } diff --git a/crates/pgls_catalog/src/catalog/mod.rs b/crates/pgls_catalog/src/catalog/mod.rs index cf5b3d5b2..f6a847352 100644 --- a/crates/pgls_catalog/src/catalog/mod.rs +++ b/crates/pgls_catalog/src/catalog/mod.rs @@ -31,17 +31,25 @@ pub use base::CatalogBase; /// `(schema, name)` of a catalog object. type Key = (String, String); -/// The state of an object that the file changed. +/// The state of an object that the file changed. Objects are shared between clones of the +/// catalog, like the checkpoints of the editor walk and savepoints, and copied when one of them +/// changes. #[derive(Debug, Clone)] enum Entry { /// Created or altered by the file. - Defined(T), + Defined(Arc), /// Dropped (or renamed away) by the file. Dropped, /// Changed by the file in a way the catalog can't follow. Unknown, } +impl Entry { + fn defined(value: T) -> Self { + Self::Defined(Arc::new(value)) + } +} + /// The catalog as seen by a statement: the database snapshot plus the changes made by the /// statements before it. #[derive(Clone)] @@ -158,7 +166,7 @@ impl Catalog { fn entry_lookup(&self, entry: &Entry) -> Lookup { match entry { - Entry::Defined(value) => Lookup::Found(value.clone()), + Entry::Defined(value) => Lookup::Found(T::clone(value)), Entry::Dropped => self.not_found(), Entry::Unknown => Lookup::Unknown, } @@ -392,7 +400,7 @@ impl CatalogView for Catalog { .values() .find_map(|entry| match entry { Entry::Defined(info) if info.id == Some(TypeId::File(*id)) => { - Some(info.clone()) + Some(TypeInfo::clone(info)) } _ => None, }) diff --git a/crates/pgls_catalog/src/catalog/overlay.rs b/crates/pgls_catalog/src/catalog/overlay.rs index 7ca62e511..94f41a361 100644 --- a/crates/pgls_catalog/src/catalog/overlay.rs +++ b/crates/pgls_catalog/src/catalog/overlay.rs @@ -1,6 +1,8 @@ //! The changes that statements are made of: defining, moving, and dropping relations, types, //! functions, and schemas in the overlay. +use std::sync::Arc; + use crate::typing::TypeId; use crate::{FunctionArgumentMode, FunctionSignature}; use pgls_query::{Node, protobuf}; @@ -82,7 +84,7 @@ impl Catalog { if relation.kind != RelationKind::Other { self.types.insert( key.clone(), - Entry::Defined(TypeInfo { + Entry::defined(TypeInfo { schema: relation.schema.clone(), name: relation.name.clone(), attributes: relation.columns.clone(), @@ -108,7 +110,7 @@ impl Catalog { ); self.ensure_array_type(&relation.schema, &relation.name, None); } - self.relations.insert(key, Entry::Defined(relation)); + self.relations.insert(key, Entry::defined(relation)); } pub(super) fn mark_relation_unknown(&mut self, key: Key) { @@ -195,26 +197,26 @@ impl Catalog { self.relations.insert(old_key.clone(), Entry::Dropped); if relation.kind != RelationKind::Other { let old_type = self.types.get(&old_key).and_then(|entry| match entry { - Entry::Defined(info) => Some(info.clone()), + Entry::Defined(info) => Some(TypeInfo::clone(info)), _ => None, }); self.types.insert(old_key.clone(), Entry::Dropped); if let Some(mut info) = old_type { let old_array = key(&info.schema, &format!("_{}", info.name)); let array_info = self.types.get(&old_array).and_then(|entry| match entry { - Entry::Defined(info) => Some(info.clone()), + Entry::Defined(info) => Some(TypeInfo::clone(info)), _ => None, }); self.types.insert(old_array, Entry::Dropped); info.schema = relation.schema.clone(); info.name = relation.name.clone(); self.types - .insert(key(&info.schema, &info.name), Entry::Defined(info.clone())); + .insert(key(&info.schema, &info.name), Entry::defined(info.clone())); if let Some(mut array) = array_info { array.schema = relation.schema.clone(); array.name = format!("_{}", relation.name); self.types - .insert(key(&array.schema, &array.name), Entry::Defined(array)); + .insert(key(&array.schema, &array.name), Entry::defined(array)); } } } @@ -259,7 +261,7 @@ impl Catalog { let array_id = self.allocate_type_id(); self.types.insert( type_key, - Entry::Defined(TypeInfo { + Entry::defined(TypeInfo { schema: schema.clone(), name: name.clone(), attributes: attributes.clone(), @@ -276,7 +278,7 @@ impl Catalog { ); self.types.insert( key(&schema, &array_name), - Entry::Defined(TypeInfo { + Entry::defined(TypeInfo { schema, name: array_name, attributes: None, @@ -313,11 +315,11 @@ impl Catalog { }) .unwrap_or_else(|| self.allocate_type_id()); if let Some(Entry::Defined(row)) = self.types.get_mut(&row_key) { - row.array = Some(array_id.clone()); + Arc::make_mut(row).array = Some(array_id.clone()); } self.types.insert( array_key, - Entry::Defined(TypeInfo { + Entry::defined(TypeInfo { schema: schema.to_owned(), name: array_name, attributes, @@ -355,7 +357,7 @@ impl Catalog { type_info.origin = Origin::File; self.types.insert(old_key, Entry::Dropped); let new_key = key(&type_info.schema, &type_info.name); - self.types.insert(new_key, Entry::Defined(type_info)); + self.types.insert(new_key, Entry::defined(type_info)); } // ----- functions ----- @@ -376,9 +378,9 @@ impl Catalog { Some(position) => overloads[position] = function, None => overloads.push(function), } - Entry::Defined(overloads) + Entry::defined(overloads) } - Lookup::Missing => Entry::Defined(vec![function]), + Lookup::Missing => Entry::defined(vec![function]), Lookup::Unknown => Entry::Unknown, }; self.functions.insert(key, entry); @@ -454,7 +456,7 @@ impl Catalog { let entry = if overloads.is_empty() { Entry::Dropped } else { - Entry::Defined(overloads) + Entry::defined(overloads) }; self.functions.insert(key, entry); Some(function) diff --git a/crates/pgls_inlay_hints/src/lib.rs b/crates/pgls_inlay_hints/src/lib.rs index 8ffd8b819..2bac7c638 100644 --- a/crates/pgls_inlay_hints/src/lib.rs +++ b/crates/pgls_inlay_hints/src/lib.rs @@ -47,8 +47,14 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { collect_references: true, }) .references; + inlay_hints_of_references(&references, params.sql) +} + +/// The inlay hints of the statement `sql` from its `references`, which +/// [`pgls_catalog::resolve::resolve`] collected. Ordered by offset. +pub fn inlay_hints_of_references(references: &[Reference], sql: &str) -> Vec { let calls: Vec<_> = references - .into_iter() + .iter() .filter_map(|reference| match reference { Reference::Function { location, @@ -76,7 +82,7 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { if calls.is_empty() { return Vec::new(); } - let Ok(scan) = pgls_query::scan(params.sql) else { + let Ok(scan) = pgls_query::scan(sql) else { return Vec::new(); }; let tokens: Vec = scan @@ -90,7 +96,7 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { let Some(signature) = function.signature.as_ref() else { continue; }; - let Some(arguments) = argument_tokens(&tokens, location) else { + let Some(arguments) = argument_tokens(&tokens, *location) else { continue; }; // An aggregate's `ORDER BY` follows its arguments, and may contain more commas. @@ -112,7 +118,7 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { else { continue; }; - if is_reference_to(params.sql, tokens, name) { + if is_reference_to(sql, tokens, name) { continue; } let Ok(offset) = u32::try_from(tokens[0].start) else { diff --git a/crates/pgls_semantic_tokens/src/lib.rs b/crates/pgls_semantic_tokens/src/lib.rs index 81aca5c13..317af7358 100644 --- a/crates/pgls_semantic_tokens/src/lib.rs +++ b/crates/pgls_semantic_tokens/src/lib.rs @@ -62,13 +62,18 @@ pub fn semantic_tokens(params: SemanticTokensParams<'_>) -> Vec { collect_references: true, }) .references; + semantic_tokens_of_references(&references, params.sql) +} - let Some(scanned) = scan(params.sql) else { +/// The semantic tokens of the statement `sql` from its `references`, which +/// [`pgls_catalog::resolve::resolve`] collected. Ordered by their start, and they don't overlap. +pub fn semantic_tokens_of_references(references: &[Reference], sql: &str) -> Vec { + let Some(scanned) = scan(sql) else { return Vec::new(); }; let mut tokens = Vec::new(); - for reference in &references { + for reference in references { if let Reference::EnumLiteral { span, type_info, .. } = reference @@ -151,8 +156,7 @@ pub fn semantic_tokens(params: SemanticTokensParams<'_>) -> Vec { } // `table.column`, but not `alias.column`. [table, _] - if relation - .is_some_and(|(_, name)| identifier(¶ms.sql[*table]) == *name) => + if relation.is_some_and(|(_, name)| identifier(&sql[*table]) == *name) => { push(Some(table), SemanticTokenKind::Class, schema); } diff --git a/crates/pgls_workspace/src/features/completions.rs b/crates/pgls_workspace/src/features/completions.rs index a5e98c885..d94e58a61 100644 --- a/crates/pgls_workspace/src/features/completions.rs +++ b/crates/pgls_workspace/src/features/completions.rs @@ -4,7 +4,7 @@ use pgls_completions::CompletionItem; use pgls_fs::PgLSPath; use pgls_text_size::TextSize; -use crate::workspace::{Document, DocumentStatement, GetCompletionsFilter, WithCSTMapper}; +use crate::workspace::{DefaultMapper, Document, DocumentStatement, WithCSTMapper}; #[derive(Debug, serde::Serialize, serde::Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -48,18 +48,14 @@ pub(crate) fn get_statement_for_completions( doc: &Document, position: TextSize, ) -> Option<(DocumentStatement, Arc)> { - let count = doc.count(); + // Only whether the document has no statement, one, or more matters. + let count = doc.iter(DefaultMapper).take(2).count(); // no arms no cookies if count == 0 { return None; } - let mut eligible_statements = doc.iter_with_filter( - WithCSTMapper, - GetCompletionsFilter { - cursor_position: position, - }, - ); + let mut eligible_statements = doc.iter_for_completions(WithCSTMapper, position); if count == 1 { eligible_statements.next() @@ -93,7 +89,7 @@ pub(crate) fn get_statement_for_completions( mod tests { use pgls_text_size::TextSize; - use crate::workspace::Document; + use crate::workspace::{DefaultMapper, Document}; use super::get_statement_for_completions; @@ -151,7 +147,7 @@ mod tests { let (doc, position) = get_doc_and_pos(sql.as_str()); // make sure these are parsed as two - assert_eq!(doc.count(), 2); + assert_eq!(doc.iter(DefaultMapper).count(), 2); assert!(get_statement_for_completions(&doc, position).is_none()); } @@ -227,4 +223,40 @@ mod tests { assert!(get_statement_for_completions(&doc, position).is_none()); } + + /// Looking only at the statements around the cursor finds the same statement as filtering + /// all statements of the document, at every position. + #[test] + fn finds_the_statement_like_a_filter_over_all_statements() { + use crate::workspace::{GetCompletionsFilter, WithCSTMapper}; + + let sql = "select 1;\nselect * from\nselect 2 select 3;\n\ncreate function f() returns int language sql as $$ select 4; select $$;\ninsert into t values (1)\n\nupdate t set a = 1;select\n"; + let doc = Document::new(sql.into(), 1); + for position in 0..=sql.len() { + let position = TextSize::new(position as u32); + let mut statements = doc.iter_with_filter( + WithCSTMapper, + GetCompletionsFilter { + cursor_position: position, + }, + ); + // The lookup before it filtered all statements. + let mut expected: Option<(crate::workspace::DocumentStatement, _)> = None; + for current in statements.by_ref() { + if expected.as_ref().is_some_and(|(previous, _)| { + current.0.range.contains(position) && !current.0.id.is_child_of(&previous.id) + }) { + expected = None; + break; + } + expected = Some(current); + } + assert_eq!( + get_statement_for_completions(&doc, position) + .map(|(stmt, _)| (stmt.id, stmt.range)), + expected.map(|(stmt, _)| (stmt.id, stmt.range)), + "at {position:?}" + ); + } + } } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 60aec0f36..dd57ba3eb 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -1,7 +1,5 @@ use std::{ - borrow::Borrow, collections::HashMap, - hash::{DefaultHasher, Hash, Hasher}, panic::RefUnwindSafe, path::{Path, PathBuf}, sync::{ @@ -26,9 +24,8 @@ use futures::{StreamExt, TryStreamExt, stream}; use pgls_analyse::{ ActionCategory, AnalysisFilter, Applicability, RuleCategories, SUPPRESSION_ACTION_CATEGORY, }; -use pgls_analyser::{LinterDiagnostic, MigrationSection, RuleFix, WalkParams}; -use pgls_catalog::resolve::{ColumnSource, Reference, ResolveParams, resolve}; -use pgls_catalog::{Catalog, CatalogBase, Origin, Session}; +use pgls_analyser::{LinterDiagnostic, MigrationSection, RuleFix}; +use pgls_catalog::CatalogBase; #[cfg(feature = "db")] use pgls_lexer::convert_to_positional_params_with_metadata; use std::borrow::Cow; @@ -64,7 +61,7 @@ use crate::{ semantic_tokens::{PullSemanticTokensParams, PullSemanticTokensResult, SemanticToken}, }, settings::{Settings, WorkspaceSettings, WorkspaceSettingsHandle, WorkspaceSettingsHandleMut}, - workspace::{AnalyserDiagnosticsMapper, WithCSTMapper}, + workspace::WithCSTMapper, }; use super::{ @@ -95,6 +92,7 @@ mod sql_function; #[cfg(feature = "db")] mod statement_execution; mod statement_identifier; +mod walk; pub struct WorkspaceServer { /// global settings object for this workspace @@ -108,23 +106,16 @@ pub struct WorkspaceServer { /// The indexed database snapshot the catalog starts from catalog_base: Mutex>>, - /// The last snapshot built for completions and hover - snapshot_before: Mutex>, - /// Incremented whenever the settings change, see [`CacheKey`]. settings_version: AtomicU64, #[cfg(feature = "db")] connection: ConnectionManager, -} -/// A database snapshot with the changes of the statements before a statement. -struct SnapshotBefore { - /// The database snapshot it is built from. - base: Arc, - /// The text that the statement sees. - key: u64, - snapshot: Arc, + /// The statements that the editor walks visited, to test that they resume from + /// checkpoints. + #[cfg(test)] + walked_statements: std::sync::atomic::AtomicUsize, } /// The `Workspace` object is long-lived, so we want it to be able to cross @@ -150,9 +141,10 @@ impl WorkspaceServer { documents: RwLock::new(HashMap::new()), snapshot: SnapshotManager::new(), catalog_base: Mutex::default(), - snapshot_before: Mutex::default(), settings_version: AtomicU64::default(), connection: ConnectionManager::new(), + #[cfg(test)] + walked_statements: Default::default(), } } @@ -164,8 +156,9 @@ impl WorkspaceServer { documents: RwLock::new(HashMap::new()), snapshot: SnapshotManager::new(), catalog_base: Mutex::default(), - snapshot_before: Mutex::default(), settings_version: AtomicU64::default(), + #[cfg(test)] + walked_statements: Default::default(), } } @@ -325,6 +318,7 @@ impl WorkspaceServer { /// ones of the file. fn refactor_code_actions( &self, + project_key: ProjectKey, settings: Option<&Settings>, doc: &Document, position: TextSize, @@ -349,248 +343,14 @@ impl WorkspaceServer { let Some(snapshot) = snapshot else { return Vec::new(); }; - let (stmt_start, references) = self.references_at(settings, doc, &snapshot, position); + let request = self.editor_request(project_key, settings, snapshot, doc); + let (stmt_start, references) = self.references_at(&request, doc, position); references .iter() .filter_map(|reference| expand_star::expand_star_action(reference, content, stmt_start)) .collect() } - /// Walks the statements of `doc` that start before `end` like the analyser does, see - /// [`pgls_analyser::walk_statements`]. `visit` gets every parsed statement, except those with - /// identifier parameters, with the catalog and the session as they are before it. `settings` - /// are the ones of the file. Returns the catalog as a statement at `end` sees it. - fn walk_statements( - &self, - settings: Option<&Settings>, - doc: &Document, - snapshot: &Arc, - end: TextSize, - mut visit: impl FnMut(&pgls_analyser::AnalysableStatement, &Catalog, &Session), - ) -> Catalog { - let (search_path, sections) = settings - .map(|settings| { - ( - settings.typecheck.search_path.clone(), - document_sections(settings, doc), - ) - }) - .unwrap_or_default(); - let mut stmts = Vec::new(); - let mut unparsable = Vec::new(); - for (stmt, _, skipped) in doc.iter(AnalyserDiagnosticsMapper) { - let Some(start) = stmt - .as_ref() - .map(|stmt| stmt.range.start()) - .or_else(|| skipped.as_ref().map(|skipped| skipped.range.start())) - else { - continue; - }; - if start >= end { - break; - } - stmts.extend(stmt); - unparsable.extend(skipped); - } - - pgls_analyser::walk_statements( - WalkParams { - stmts: &stmts, - unparsable: &unparsable, - catalog_base: Some(self.catalog_base(snapshot)), - search_path: &snapshot.expand_search_path(&search_path), - sections: §ions, - end, - // Editor features show what names most likely refer to, so a `DO` block or a - // statement that doesn't parse doesn't turn them off for the rest of the file. - best_effort: true, - }, - |stmt, file_context| { - // The names of statements with identifier parameters like `:name` are - // placeholders. - if !stmt.has_identifier_parameters { - visit(stmt, file_context.catalog(), file_context.session()); - } - }, - ) - } - - /// For hover on the statement `id` at `range`: the references whose span contains - /// `position`, and the database snapshot as the statement sees it. - fn hover_references( - &self, - settings: Option<&Settings>, - doc: &Document, - id: &StatementId, - range: TextRange, - position: TextSize, - snapshot: Arc, - ) -> (Vec, Arc) { - let search_path = settings - .map(|settings| settings.typecheck.search_path.clone()) - .unwrap_or_default(); - let mut found = None; - let catalog = self.walk_statements( - settings, - doc, - &snapshot, - position + TextSize::from(1), - |stmt, catalog, session| { - let Some(sql) = stmt.sql.as_deref() else { - return; - }; - if stmt.range != range { - return; - } - let position = stmt.statement_offset(position); - let (references, others): (Vec<_>, Vec<_>) = resolve(ResolveParams { - stmt: &stmt.root, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - sql: Some(sql), - collect_references: true, - }) - .references - .into_iter() - .partition(|reference| { - reference - .span() - .is_some_and(|span| span.contains_inclusive(position)) - }); - // Objects that the file didn't change are the same in the database snapshot. - let snapshot = if in_database(&references, &others) { - Arc::clone(&snapshot) - } else { - self.snapshot_before( - doc, - id, - range, - Arc::clone(&snapshot), - &search_path, - || catalog, - ) - }; - found = Some((references, snapshot)); - }, - ); - // A statement that doesn't parse is not visited. It changes nothing, so it sees the - // catalog after the walk. - found.unwrap_or_else(|| { - ( - Vec::new(), - self.snapshot_before(doc, id, range, snapshot, &search_path, || catalog), - ) - }) - } - - /// The references whose span contains `position`, in the innermost statement at `position` - /// (the body of a SQL function rather than its `CREATE FUNCTION`). Their spans are relative - /// to the start of that statement, which is returned with them. Empty if the statement's - /// text isn't the source text, like a function body with escaped quotes, since the spans - /// would then not point into the document. - fn references_at( - &self, - settings: Option<&Settings>, - doc: &Document, - snapshot: &Arc, - position: TextSize, - ) -> (TextSize, Vec) { - let mut innermost: Option<(TextRange, bool, Vec)> = None; - self.walk_statements( - settings, - doc, - snapshot, - position + TextSize::from(1), - |stmt, catalog, session| { - let Some(sql) = stmt.sql.as_deref() else { - return; - }; - if !stmt.range.contains(position) - || innermost - .as_ref() - .is_some_and(|(range, ..)| !range.contains_range(stmt.range)) - { - return; - } - let references = resolve(ResolveParams { - stmt: &stmt.root, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - sql: Some(sql), - collect_references: true, - }) - .references; - innermost = Some((stmt.range, stmt.offsets.is_verbatim(), references)); - }, - ); - let Some((range, true, references)) = innermost else { - return (TextSize::default(), Vec::new()); - }; - let position = position - range.start(); - let references = references - .into_iter() - .filter(|reference| { - reference - .span() - .is_some_and(|span| span.contains_inclusive(position)) - }) - .collect(); - (range.start(), references) - } - - /// The indexed database snapshot, built once per snapshot. - fn catalog_base(&self, snapshot: &Arc) -> Arc { - let mut cached = self.catalog_base.lock().unwrap(); - if let Some(base) = cached.as_ref().filter(|base| base.is_built_from(snapshot)) { - return Arc::clone(base); - } - let base = Arc::new(CatalogBase::new(Arc::clone(snapshot))); - *cached = Some(Arc::clone(&base)); - base - } - - /// The database snapshot as the statement `id` at `range` sees it: with the objects that the - /// statements before it create, change, or drop. `catalog` returns the catalog that the - /// statement sees, and is only called when the snapshot is not cached. - fn snapshot_before>( - &self, - doc: &Document, - id: &StatementId, - range: TextRange, - snapshot: Arc, - search_path: &[String], - catalog: impl FnOnce() -> C, - ) -> Arc { - // The statement sees the text before it, and the body of a SQL function also sees its - // `CREATE FUNCTION`. Files with other settings can see another search path. - let mut hasher = DefaultHasher::new(); - doc.get_document_content()[..usize::from(range.start())].hash(&mut hasher); - id.parent_content().hash(&mut hasher); - search_path.hash(&mut hasher); - let key = hasher.finish(); - - let mut cached = self.snapshot_before.lock().unwrap(); - if let Some(cached) = cached.as_ref() - && cached.key == key - && Arc::ptr_eq(&cached.base, &snapshot) - { - return Arc::clone(&cached.snapshot); - } - - let result = catalog() - .borrow() - .snapshot() - .unwrap_or_else(|| Arc::clone(&snapshot)); - *cached = Some(SnapshotBefore { - base: snapshot, - key, - snapshot: Arc::clone(&result), - }); - result - } - /// Diagnostics from the database: EXPLAIN for top-level statements whose references are all /// unchanged database objects, and plpgsql_check. Also returns the ranges of the /// statements that plpgsql_check checked. @@ -812,8 +572,6 @@ impl Workspace for WorkspaceServer { params.vcs_base_path, params.gitignore_matches.as_slice(), )?; - // The search path and the migration sections change what a statement sees. - *self.snapshot_before.lock().unwrap() = None; Ok(()) } @@ -882,8 +640,6 @@ impl Workspace for WorkspaceServer { self.workspaces_mut() .as_mut() .set_nested_settings(params.project_key, nested); - // The search path and the migration sections change what a statement sees. - *self.snapshot_before.lock().unwrap() = None; Ok(result) } @@ -919,16 +675,17 @@ impl Workspace for WorkspaceServer { fn change_file(&self, params: super::ChangeFileParams) -> Result<(), WorkspaceError> { let mut documents = self.documents.write().unwrap(); - match documents.entry(params.path.clone()) { - std::collections::hash_map::Entry::Occupied(mut entry) => { - entry - .get_mut() - .update_content(params.content, params.version); - } + let stale = match documents.entry(params.path.clone()) { + std::collections::hash_map::Entry::Occupied(mut entry) => entry + .get_mut() + .update_content(params.content, params.version), std::collections::hash_map::Entry::Vacant(entry) => { entry.insert(Document::new(params.content, params.version)); + return Ok(()); } - } + }; + drop(documents); + stale.drop_in_background(); Ok(()) } @@ -988,7 +745,12 @@ impl Workspace for WorkspaceServer { actions.splice( 0..0, - self.refactor_code_actions(settings, parser, params.cursor_position), + self.refactor_code_actions( + params.project_key, + settings, + parser, + params.cursor_position, + ), ); if let Some(settings) = settings { @@ -1517,14 +1279,10 @@ impl Workspace for WorkspaceServer { .ok_or(WorkspaceError::not_found())?; let settings = self.workspaces(); - let search_path = settings - .settings_for(params.project_key, params.path.as_path()) - .map(|settings| settings.typecheck.search_path.clone()) - .unwrap_or_default(); + let settings = settings.settings_for(params.project_key, params.path.as_path()); #[cfg(feature = "db")] let Some(snapshot) = settings - .settings_for(params.project_key, params.path.as_path()) .and_then(|settings| { self.connection .with_pool(&settings.db, |pool| self.snapshot.load(pool)) @@ -1548,23 +1306,9 @@ impl Workspace for WorkspaceServer { } Some((stmt, cst)) => { let position = stmt.statement_offset(params.position); - let range = stmt.range; - let snapshot = self.snapshot_before( - parsed_doc, - &stmt.id, - range, - Arc::clone(&snapshot), - &search_path, - || { - self.walk_statements( - settings.settings_for(params.project_key, params.path.as_path()), - parsed_doc, - &snapshot, - range.start(), - |_, _, _| {}, - ) - }, - ); + let request = + self.editor_request(params.project_key, settings, snapshot, parsed_doc); + let snapshot = self.snapshot_before(&request, parsed_doc, &stmt); let items = pgls_completions::complete(pgls_completions::CompletionParams { position, @@ -1614,14 +1358,9 @@ impl Workspace for WorkspaceServer { match doc.iter_at(WithCSTMapper, params.position).last() { Some((stmt, ts_tree)) => { let position_in_stmt = stmt.statement_offset(params.position); - let (references, snapshot) = self.hover_references( - settings, - doc, - &stmt.id, - stmt.range, - params.position, - snapshot, - ); + let request = self.editor_request(params.project_key, settings, snapshot, doc); + let (references, snapshot) = + self.hover_references(&request, doc, &stmt, params.position); let markdown_blocks = pgls_hover::on_hover(pgls_hover::OnHoverParams { ts_tree: &ts_tree, @@ -1676,36 +1415,28 @@ impl Workspace for WorkspaceServer { return Ok(PullInlayHintsResult::default()); }; + let request = self.editor_request(params.project_key, settings, snapshot, doc); let mut hints = Vec::new(); - self.walk_statements( - settings, - doc, - &snapshot, - params.range.end(), - |stmt, catalog, session| { + for statements in self.resolved_statements(&request, doc, params.range) { + for resolved in statements.iter() { + let stmt = &resolved.stmt; let Some(sql) = stmt.sql.as_deref() else { - return; + continue; }; if stmt.range.intersect(params.range).is_none() { - return; + continue; } hints.extend( - pgls_inlay_hints::inlay_hints(pgls_inlay_hints::InlayHintsParams { - stmt: &stmt.root, - sql, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - }) - .into_iter() - .map(|hint| InlayHint { - offset: stmt.document_offset(hint.offset), - label: hint.label, - kind: hint.kind.into(), - }), + pgls_inlay_hints::inlay_hints_of_references(&resolved.references, sql) + .into_iter() + .map(|hint| InlayHint { + offset: stmt.document_offset(hint.offset), + label: hint.label, + kind: hint.kind.into(), + }), ); - }, - ); + } + } hints.sort_by_key(|hint| hint.offset); Ok(PullInlayHintsResult { hints }) @@ -1752,38 +1483,28 @@ impl Workspace for WorkspaceServer { let range = params .range .unwrap_or_else(|| TextRange::up_to(TextSize::of(doc.get_document_content()))); + let request = self.editor_request(params.project_key, settings, snapshot, doc); let mut tokens = Vec::new(); - self.walk_statements( - settings, - doc, - &snapshot, - range.end(), - |stmt, catalog, session| { + for statements in self.resolved_statements(&request, doc, range) { + for resolved in statements.iter() { + let stmt = &resolved.stmt; let Some(sql) = stmt.sql.as_deref() else { - return; + continue; }; if stmt.range.intersect(range).is_none() { - return; + continue; } tokens.extend( - pgls_semantic_tokens::semantic_tokens( - pgls_semantic_tokens::SemanticTokensParams { - stmt: &stmt.root, - sql, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - }, - ) - .into_iter() - .map(|token| SemanticToken { - range: stmt.document_range(token.range), - kind: token.kind.into(), - default_library: token.default_library, - }), + pgls_semantic_tokens::semantic_tokens_of_references(&resolved.references, sql) + .into_iter() + .map(|token| SemanticToken { + range: stmt.document_range(token.range), + kind: token.kind.into(), + default_library: token.default_library, + }), ); - }, - ); + } + } // The body of a SQL function is a statement of its own, inside its `CREATE FUNCTION`. tokens.sort_by_key(|token| (token.range.start(), token.range.end())); @@ -1799,41 +1520,6 @@ impl Workspace for WorkspaceServer { } } -/// Whether hover finds what `references` resolve to in the database snapshot: objects that the -/// file didn't change, and columns of CTEs and subqueries. `others` are the other references of -/// the statement. -fn in_database(references: &[Reference], others: &[Reference]) -> bool { - // The statement references each relation that it takes columns from. - let unchanged_relation = |schema: &str, name: &str| { - references.iter().chain(others).any(|reference| { - matches!( - reference, - Reference::Relation { relation, .. } - if relation.schema == schema - && relation.name == name - && relation.origin == Origin::Database - ) - }) - }; - !references.is_empty() - && references.iter().all(|reference| match reference { - Reference::Relation { relation, .. } => relation.origin == Origin::Database, - Reference::Function { function, .. } => function.origin == Origin::Database, - Reference::Column { - source: ColumnSource::Relation { schema, name }, - .. - } => unchanged_relation(schema, name), - Reference::Column { .. } => true, - Reference::Type { type_info, .. } | Reference::EnumLiteral { type_info, .. } => { - type_info.origin == Origin::Database - } - // Parameters belong to a function of the file. - Reference::Parameter { .. } => false, - // The columns of a `*` can come from objects that the file changed. - Reference::Star { .. } => false, - }) -} - /// The migration sections of `doc`, like the up and down migrations of dbmate and goose. fn document_sections(settings: &Settings, doc: &Document) -> Vec { settings diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 9d794fec1..dc250d082 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -4709,3 +4709,475 @@ async fn lint_hover_and_completions_see_the_same_catalog(test_db: PgPool) { assert!(!completions.contains(&"id".to_string()), "{completions:?}"); assert_eq!(hovers, [""]); } + +/// Hover and completions resume the walk over the statements from a checkpoint before the +/// statement, as long as the text before the checkpoint doesn't change. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn walks_resume_from_checkpoints_before_the_edit(test_db: PgPool) { + use std::sync::atomic::Ordering; + + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + migrations: Some(PartialMigrationsConfiguration { + section_markers: Some(SectionMarkers::Dbmate), + ..Default::default() + }), + ..Default::default() + }); + let (workspace, project_key) = + get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + let path = PgLSPath::new("20250301120000_migration.sql"); + + let tables: String = (0..200) + .map(|i| format!("create table checkpoint_t{i} (id int);\n")) + .collect(); + let file = |early_columns: &str, before_last: &str, last: &str| { + format!("create table checkpoint_early ({early_columns});\n{tables}{before_last}{last}\n") + }; + let version = std::cell::Cell::new(1); + let change = |content: &str| { + version.set(version.get() + 1); + workspace + .change_file(crate::workspace::ChangeFileParams { + project_key, + path: path.clone(), + content: content.into(), + version: version.get(), + }) + .expect("Unable to change test file"); + }; + // The statements that the request walked, and its result. + let walked = |request: &dyn Fn() -> Vec| { + let before = workspace.walked_statements.load(Ordering::Relaxed); + let result = request(); + ( + workspace.walked_statements.load(Ordering::Relaxed) - before, + result, + ) + }; + let hover = |content: &str, needle: &str| { + let offset = content.rfind(needle).expect("needle in content") + 2; + workspace + .on_hover(OnHoverParams { + project_key, + path: path.clone(), + position: u32::try_from(offset).unwrap().into(), + }) + .expect("Unable to hover") + .markdown_blocks + }; + let complete = |content: &str, needle: &str| { + let offset = content.rfind(needle).expect("needle in content") + needle.len(); + workspace + .get_completions(crate::workspace::GetCompletionsParams { + project_key, + path: path.clone(), + position: u32::try_from(offset).unwrap().into(), + }) + .expect("Unable to request completions") + .items + .into_iter() + .map(|item| item.label) + .collect::>() + }; + + let content = file("id int", "", "select id from checkpoint_early;"); + workspace + .open_file(OpenFileParams { + project_key, + path: path.clone(), + content: content.clone(), + version: 1, + }) + .expect("Unable to open test file"); + let (count, blocks) = walked(&|| hover(&content, "checkpoint_early;")); + assert!(count > 200, "{count}"); + assert!(blocks[0].contains("- id: int4"), "{blocks:?}"); + // A repeated request finds the references and the snapshot in the cache. + let (count, blocks) = walked(&|| hover(&content, "checkpoint_early;")); + assert_eq!(count, 0); + assert!(blocks[0].contains("- id: int4"), "{blocks:?}"); + + // An edit in the hovered statement keeps the checkpoint before it. + let content = file("id int", "", "select id from checkpoint_early;"); + change(&content); + let (count, blocks) = walked(&|| hover(&content, "checkpoint_early;")); + assert_eq!(count, 1); + assert!(blocks[0].contains("- id: int4"), "{blocks:?}"); + + // So does a new statement after it. Completions walk to the start of their statement. + let content = file( + "id int", + "select id from checkpoint_early;\n", + "select from checkpoint_early;", + ); + change(&content); + let (count, labels) = walked(&|| complete(&content, "select ")); + assert_eq!(count, 1); + assert!(labels.contains(&"id".to_string()), "{labels:?}"); + + // An edit before the checkpoints drops them: the statements after it see the new column. + let content = file( + "id int, extra int", + "select id from checkpoint_early;\n", + "select from checkpoint_early;", + ); + change(&content); + let (count, labels) = walked(&|| complete(&content, "select ")); + assert!(count > 200, "{count}"); + assert!(labels.contains(&"extra".to_string()), "{labels:?}"); + // Hover on the same statement resumes from the checkpoint that completions recorded. + let (count, blocks) = walked(&|| hover(&content, "checkpoint_early;")); + assert_eq!(count, 1); + assert!(blocks[0].contains("- extra: int4"), "{blocks:?}"); + + // The walk also records a checkpoint at the statement where the last change starts, so the + // next keystroke there walks from it. + let content = file( + "id int, extra int", + "select id from checkpoint_t1;\n", + "select from checkpoint_early;", + ); + change(&content); + let (count, _) = walked(&|| hover(&content, "checkpoint_early;")); + assert!(count > 200, "{count}"); + let content = file( + "id int, extra int", + "select id from checkpoint_t12;\n", + "select from checkpoint_early;", + ); + change(&content); + let (count, blocks) = walked(&|| hover(&content, "checkpoint_early;")); + assert_eq!(count, 2); + assert!(blocks[0].contains("- extra: int4"), "{blocks:?}"); + + // A section marker right before the statement of a checkpoint puts it in another section, + // so the checkpoint doesn't fit anymore. The walk resumes from the one before. + let marked = content.replace( + "select from checkpoint_early;", + "-- migrate:down\nselect from checkpoint_early;", + ); + change(&marked); + let (count, _) = walked(&|| hover(&marked, "checkpoint_early;")); + assert_eq!(count, 2); + + let documents = workspace.documents.read().unwrap(); + assert!(documents.get(&path).unwrap().cache().walk_size().1 <= 4); +} + +/// After edits anywhere in the file, hover and inlay hints that resume from checkpoints give +/// the same results as on a document that is walked from its start. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn walks_from_checkpoints_match_full_walks(test_db: PgPool) { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + migrations: Some(PartialMigrationsConfiguration { + section_markers: Some(SectionMarkers::Dbmate), + ..Default::default() + }), + ..Default::default() + }); + let (workspace, project_key) = + get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + let edited = PgLSPath::new("20250301120000_edited.sql"); + let fresh = PgLSPath::new("20250301120000_fresh.sql"); + + let mut statements: Vec = vec!["-- migrate:up".into()]; + for i in 0..30 { + statements.push(format!("create table eq_t{i} (id int, name text);")); + statements.push(format!( + "create function eq_f{i}(a int, b text) returns int language sql as 'select a';" + )); + statements.push(format!("select eq_f{i}(id, name) from eq_t{i};")); + if i % 10 == 9 { + statements.push("-- migrate:down".into()); + statements.push(format!("alter table eq_t{i} add column down_only int;")); + statements.push("create tabel eq_unparsed (id int);".into()); + statements.push("-- migrate:up".into()); + } + } + // Statements at the end that use objects of the whole file. + for i in [1, 9, 10, 25] { + statements.push(format!("select eq_f{i}(1, 'x'), * from eq_t{i};")); + } + let render = |statements: &[String]| statements.join("\n") + "\n"; + + // Requests on the statements from `from` on. + let requests = |path: &PgLSPath, content: &str, from: usize| { + let mut results = Vec::new(); + for (offset, _) in content + .rmatch_indices("from eq_t") + .filter(|(o, _)| *o >= from) + { + // On the table name, from the end of the file, so that the requests don't only + // resume from the checkpoints that the requests before them recorded. + let blocks = workspace + .on_hover(OnHoverParams { + project_key, + path: path.clone(), + position: u32::try_from(offset + 7).unwrap().into(), + }) + .expect("Unable to hover") + .markdown_blocks; + results.push(blocks.join("\n")); + } + let hints = workspace + .pull_inlay_hints(PullInlayHintsParams { + project_key, + path: path.clone(), + range: TextRange::new( + u32::try_from(from).unwrap().into(), + u32::try_from(content.len()).unwrap().into(), + ), + }) + .expect("Unable to pull inlay hints") + .hints; + results.extend( + hints + .iter() + .map(|hint| format!("{}@{}", hint.label, u32::from(hint.offset))), + ); + let tokens = workspace + .pull_semantic_tokens(PullSemanticTokensParams { + project_key, + path: path.clone(), + range: Some(TextRange::new( + u32::try_from(from).unwrap().into(), + u32::try_from(content.len()).unwrap().into(), + )), + }) + .expect("Unable to pull semantic tokens") + .tokens; + results.extend( + tokens + .iter() + .map(|token| format!("{:?}@{:?}", token.kind, token.range)), + ); + results + }; + + let content = render(&statements); + workspace + .open_file(OpenFileParams { + project_key, + path: edited.clone(), + content: content.clone(), + version: 1, + }) + .expect("Unable to open test file"); + let results = requests(&edited, &content, content.len() / 2); + assert!( + results.iter().any(|result| result.starts_with("a:")), + "{results:?}" + ); + assert!( + results.iter().any(|result| result.contains("- name: text")), + "{results:?}" + ); + + // Edits early, in the middle and late in the file, and twice in a row at the same place, with + // requests at and after them. + let edits: [(usize, &str); 8] = [ + (4, "create table eq_t1 (id int, name text, extra int);"), + ( + 5, + "create function eq_f1(x int, y text) returns int language sql as 'select x';", + ), + ( + 5, + "create function eq_f1(xx int, y text) returns int language sql as 'select xx';", + ), + ( + 36, + "create function eq_f10(c int, d text) returns int language sql as 'select c';", + ), + (32, "alter table eq_t9 drop column name;"), + (33, "create table eq_unparsed (id int);"), + (84, "create table eq_t25 (id int);"), + (84, "create table eq_t25 (id bigint);"), + ]; + for (version, (i, statement)) in edits.into_iter().enumerate() { + // Each edit changes a statement of the same kind. + assert_eq!(statements[i].split(' ').next(), statement.split(' ').next(),); + statements[i] = statement.to_string(); + let content = render(&statements); + workspace + .change_file(crate::workspace::ChangeFileParams { + project_key, + path: edited.clone(), + content: content.clone(), + version: i32::try_from(version).unwrap() + 2, + }) + .expect("Unable to change test file"); + workspace + .open_file(OpenFileParams { + project_key, + path: fresh.clone(), + content: content.clone(), + version: 1, + }) + .expect("Unable to open test file"); + + let edit_offset = content.find(statement).unwrap(); + for from in [edit_offset, content.len() / 2, content.len() * 3 / 4] { + assert_eq!( + requests(&edited, &content, from), + requests(&fresh, &content, from), + "after editing statement {i}, from {from}" + ); + // The fresh document walks from its start for every request. + workspace + .close_file(crate::workspace::CloseFileParams { + project_key, + path: fresh.clone(), + }) + .expect("Unable to close test file"); + workspace + .open_file(OpenFileParams { + project_key, + path: fresh.clone(), + content: content.clone(), + version: 1, + }) + .expect("Unable to open test file"); + } + workspace + .close_file(crate::workspace::CloseFileParams { + project_key, + path: fresh.clone(), + }) + .expect("Unable to close test file"); + } +} + +/// Whether a statement runs in an implicit transaction block depends on the query it is in, +/// also on its statements after it. Cached walks see the blocks like a walk from the start. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn cached_walks_see_the_implicit_blocks_of_queries(test_db: PgPool) { + test_db + .execute( + "create schema alt; create table public.t (pub_col int); create table alt.t (alt_col int);", + ) + .await + .expect("setup sql failed"); + + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + migrations: Some(PartialMigrationsConfiguration { + section_markers: Some(SectionMarkers::Goose), + ..Default::default() + }), + ..Default::default() + }); + let (workspace, project_key) = + get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + let edited = PgLSPath::new("20250301120000_edited.sql"); + let fresh = PgLSPath::new("20250301120000_fresh.sql"); + + let open = |path: &PgLSPath, content: &str| { + workspace + .open_file(OpenFileParams { + project_key, + path: path.clone(), + content: content.into(), + version: 1, + }) + .expect("Unable to open test file"); + }; + // Hover on the table of the `nth` `from t`. + let hover = |path: &PgLSPath, content: &str, nth: usize| { + let (offset, _) = content.match_indices("from t").nth(nth).expect("needle"); + workspace + .on_hover(OnHoverParams { + project_key, + path: path.clone(), + position: u32::try_from(offset + 5).unwrap().into(), + }) + .expect("Unable to hover") + .markdown_blocks + .join("\n") + }; + let fresh_hover = |content: &str, nth: usize| { + open(&fresh, content); + let blocks = hover(&fresh, content, nth); + workspace + .close_file(crate::workspace::CloseFileParams { + project_key, + path: fresh.clone(), + }) + .expect("Unable to close test file"); + blocks + }; + + // `SET LOCAL` only lasts until the end of the implicit block of its query. Without + // `StatementEnd`, the statements are sent one by one, and it does nothing. + let unfinished = "-- +goose NO TRANSACTION\n-- +goose Up\n-- +goose StatementBegin\nset local search_path = alt;\nselect * from t;\n"; + open(&edited, unfinished); + let blocks = hover(&edited, unfinished, 0); + assert!(blocks.contains("pub_col"), "{blocks}"); + // Ending the block after the statement puts it in the implicit block. + let finished = format!("{unfinished}-- +goose StatementEnd\nselect * from t;\n"); + workspace + .change_file(crate::workspace::ChangeFileParams { + project_key, + path: edited.clone(), + content: finished.clone(), + version: 2, + }) + .expect("Unable to change test file"); + let blocks = hover(&edited, &finished, 0); + assert!(blocks.contains("alt_col"), "{blocks}"); + assert_eq!(blocks, fresh_hover(&finished, 0)); + + // The hover in the block recorded a checkpoint in it. The walk to the statement after the + // block resumes from it, and counts the statements of the block before the checkpoint: the + // block ends before that statement. + let blocks = hover(&edited, &finished, 1); + assert!(blocks.contains("pub_col"), "{blocks}"); + assert_eq!(blocks, fresh_hover(&finished, 1)); + + // That hover recorded a checkpoint before the statement after the block, where the block + // hasn't ended yet. After an edit in the statement, the walk resumes from it. + let edited_after = format!("{unfinished}-- +goose StatementEnd\nselect * from t;\n"); + workspace + .change_file(crate::workspace::ChangeFileParams { + project_key, + path: edited.clone(), + content: edited_after.clone(), + version: 3, + }) + .expect("Unable to change test file"); + let blocks = hover(&edited, &edited_after, 1); + assert!(blocks.contains("pub_col"), "{blocks}"); + assert_eq!(blocks, fresh_hover(&edited_after, 1)); +} diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 377e39a7e..2712286ce 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -7,7 +7,10 @@ use pgls_suppressions::Suppressions; use pgls_text_size::{TextRange, TextSize}; use rustc_hash::FxHashSet; -use super::{document_cache::DocumentCache, statement_identifier::StatementId}; +use super::{ + document_cache::{ContentChange, DocumentCache, StaleEntries}, + statement_identifier::StatementId, +}; pub struct Document { content: String, @@ -16,6 +19,8 @@ pub struct Document { diagnostics: Vec, cache: DocumentCache, suppressions: Suppressions, + /// Where the last change of the content starts. + edited_at: Option, } impl Document { @@ -30,15 +35,36 @@ impl Document { version, cache: DocumentCache::new(), suppressions, + edited_at: None, } } - pub fn update_content(&mut self, content: String, version: i32) { + /// Replaces the content of the document. Returns what the change invalidated in its cache, + /// for the caller to drop, see [`StaleEntries::drop_in_background`]. + pub fn update_content(&mut self, content: String, version: i32) -> StaleEntries { + let unchanged = common_prefix(&self.content, &content); self.content = content; self.version = version; let (ranges, diagnostics) = split_with_diagnostics(&self.content, None); + // The statements before the change keep what the walks cached for them. The splitter + // looks ahead, so a statement before the change can still end elsewhere. + let first_changed_statement = self + .ranges + .iter() + .zip(&ranges) + .find(|(old, new)| old != new || old.end() > unchanged) + .map_or_else( + || { + self.ranges + .get(ranges.len()) + .map_or(TextSize::from(u32::MAX), |range| range.start()) + }, + |(old, _)| old.start(), + ); + self.edited_at = Some(unchanged); + self.suppressions = Suppressions::from(self.content.as_str()).with_statements(ranges.clone()); self.ranges = ranges; @@ -53,7 +79,15 @@ impl Document { let is_present = |id: &StatementId| statements.contains(id.parent_content().unwrap_or(id.content())); self.cache.retain_statements(is_present); - self.cache.content_changed(); + self.cache.content_changed(ContentChange { + unchanged, + first_changed_statement, + }) + } + + /// Where the last change of the content starts. + pub fn edited_at(&self) -> Option { + self.edited_at } /// What is cached for the document. @@ -121,8 +155,43 @@ impl Document { } } - pub fn count(&self) -> usize { - self.iter(DefaultMapper).count() + /// The statements from the first top-level statement that starts at or after `offset` on. + pub fn iter_from<'a, M>(&'a self, mapper: M, offset: TextSize) -> ParseIterator<'a, M, NoFilter> + where + M: StatementMapper<'a>, + { + let first = self.ranges.partition_point(|range| range.start() < offset); + ParseIterator { + ranges: self.ranges[first..].iter(), + ..ParseIterator::new(self, mapper, NoFilter) + } + } + + /// The statements that [`GetCompletionsFilter`] accepts for `position`, without looking at + /// the statements far from it. + pub fn iter_for_completions<'a, M>( + &'a self, + mapper: M, + position: TextSize, + ) -> ParseIterator<'a, M, GetCompletionsFilter> + where + M: StatementMapper<'a>, + { + // The filter accepts a statement up to two characters after its end. The bodies of SQL + // functions are inside their top-level statement. + let ranges = &self.ranges; + let first = ranges.partition_point(|range| range.end() + TextSize::from(2) <= position); + let last = first + ranges[first..].partition_point(|range| range.start() <= position); + ParseIterator { + ranges: ranges[first..last].iter(), + ..ParseIterator::new( + self, + mapper, + GetCompletionsFilter { + cursor_position: position, + }, + ) + } } } @@ -445,6 +514,28 @@ impl StatementFilter<'_> for IdFilter { } } +/// The length of the longest common prefix of `a` and `b`, on a character boundary. +fn common_prefix(a: &str, b: &str) -> TextSize { + /// Compares this many bytes at a time, which is faster than comparing each byte. + const BLOCK_BYTES: usize = 64; + let (a, b) = (a.as_bytes(), b.as_bytes()); + let len = a.len().min(b.len()); + let mut equal = 0; + while equal + BLOCK_BYTES <= len + && a[equal..equal + BLOCK_BYTES] == b[equal..equal + BLOCK_BYTES] + { + equal += BLOCK_BYTES; + } + while equal < len && a[equal] == b[equal] { + equal += 1; + } + // A change inside a character starts at its first byte. + while equal > 0 && (equal < a.len() && (a[equal] & 0xC0) == 0x80) { + equal -= 1; + } + TextSize::try_from(equal).unwrap_or_default() +} + /// Helper function that wraps the statement splitter and returns the ranges with unified /// diagnostics pub(crate) fn split_with_diagnostics( @@ -839,7 +930,8 @@ $$ LANGUAGE plpgsql;"; let second_ast = d.cache.parse(&second); let first_tree = d.cache.tree(&first); - d.update_content("select 1;\nselect 3;".to_string(), 2); + d.update_content("select 1;\nselect 3;".to_string(), 2) + .drop_in_background(); assert!(Arc::ptr_eq(&first_ast, &d.cache.parse(&first))); assert!(Arc::ptr_eq(&first_tree, &d.cache.tree(&first))); diff --git a/crates/pgls_workspace/src/workspace/server/document_cache.rs b/crates/pgls_workspace/src/workspace/server/document_cache.rs index 9c119a995..f64cfb643 100644 --- a/crates/pgls_workspace/src/workspace/server/document_cache.rs +++ b/crates/pgls_workspace/src/workspace/server/document_cache.rs @@ -15,10 +15,53 @@ use crate::workspace::ProjectKey; mod parse; mod tree; +mod walk; pub(crate) use parse::ParseResult; use parse::PgQueryStore; use tree::TreeSitterStore; +pub(crate) use walk::{ContentChange, MigrationContext, ResolvedStatement, WalkCache}; + +/// What an edit invalidated in a document's cache. On a long file, the analysis, the copies of +/// the catalog in the walk checkpoints and the snapshot hold every object the file defines, and +/// dropping them takes tens of milliseconds, so [`StaleEntries::drop_in_background`] drops them +/// off the path of the edit. +#[must_use] +pub(crate) struct StaleEntries { + analysis: Option>, + walk: WalkCache, +} + +/// The stack of the thread that drops [`StaleEntries`]. Dropping a statement's tree is recursive +/// like parsing it, so this is the stack that the CLI gives its workers for parsing. +#[cfg(not(target_family = "wasm"))] +const DROP_THREAD_STACK_BYTES: usize = 8 * 1024 * 1024; + +impl StaleEntries { + /// Drops the entries on a thread of its own, so that the edit doesn't wait for it. If the + /// thread can't be started, and in WebAssembly, which has no threads, they are dropped right + /// away. + pub(crate) fn drop_in_background(self) { + if self.analysis.is_none() && self.walk.is_empty() { + return; + } + #[cfg(target_family = "wasm")] + drop(self); + #[cfg(not(target_family = "wasm"))] + self.drop_on_thread(); + } + + #[cfg(not(target_family = "wasm"))] + fn drop_on_thread(self) { + let spawned = std::thread::Builder::new() + .name("pgls-cache-drop".into()) + .stack_size(DROP_THREAD_STACK_BYTES) + .spawn(move || drop(self)); + if let Err(error) = spawned { + tracing::debug!("Dropped the stale cache entries of an edit in place: {error}"); + } + } +} /// What the caches of a document depend on besides its text. #[derive(Clone)] @@ -71,6 +114,7 @@ pub(crate) struct DocumentCache { struct KeyedCaches { key: Option, analysis: Option>, + walk: WalkCache, } impl DocumentCache { @@ -104,6 +148,12 @@ impl DocumentCache { self.keyed.lock().unwrap().analysis.clone() } + /// The number of resolved statements and walk checkpoints, whatever their key. + #[cfg(test)] + pub(crate) fn walk_size(&self) -> (usize, usize) { + self.keyed.lock().unwrap().walk.len() + } + /// The caches for `key`, emptied first if they were for another key. fn keyed(&self, key: &CacheKey) -> std::sync::MutexGuard<'_, KeyedCaches> { let mut keyed = self.keyed.lock().unwrap(); @@ -116,10 +166,35 @@ impl DocumentCache { keyed } - /// Drops what the new content of the document invalidates, see [`Self::retain_statements`] - /// for the statements. - pub(crate) fn content_changed(&mut self) { - self.keyed.get_mut().unwrap().analysis = None; + /// Runs `f` on what the walks over the document left for `key`, without what was computed + /// in another migration `context`, see [`WalkCache::set_context`]. It is not held while + /// walking, so that requests on the same document don't wait for each other. + pub(crate) fn with_walk( + &self, + key: &CacheKey, + context: &MigrationContext, + f: impl FnOnce(&mut WalkCache) -> R, + ) -> R { + let mut keyed = self.keyed(key); + let stale = keyed.walk.set_context(context); + let result = f(&mut keyed.walk); + drop(keyed); + StaleEntries { + analysis: None, + walk: stale, + } + .drop_in_background(); + result + } + + /// Takes out what `change` invalidates, see [`WalkCache::content_changed`], and the + /// analysis. [`Self::retain_statements`] drops what was cached for the statements. + pub(crate) fn content_changed(&mut self, change: ContentChange) -> StaleEntries { + let keyed = self.keyed.get_mut().unwrap(); + StaleEntries { + analysis: keyed.analysis.take(), + walk: keyed.walk.content_changed(change), + } } /// The tree of the statement `id`, or why it doesn't parse. diff --git a/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs b/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs new file mode 100644 index 000000000..7035ec2c9 --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/document_cache/walk.rs @@ -0,0 +1,255 @@ +//! What the walk over a document leaves for the next requests: states to resume it from, the +//! references of the statements it resolved, and the database snapshot as one statement sees +//! it. +//! +//! All of it depends on the text before its top-level statement, and the resolved references +//! also on the statement itself. An edit keeps what comes before it, see [`ContentChange`]. It +//! also depends on the migration sections before it, and on whether the queries it is in run in +//! an implicit transaction block, which text after it can change, see [`MigrationContext`]. + +use std::{collections::BTreeMap, sync::Arc}; + +use pgls_analyser::{AnalysableStatement, Direction, MigrationSection, WalkState}; +use pgls_catalog::{Snapshot, resolve::Reference}; +use pgls_text_size::{TextRange, TextSize}; + +/// The states a document keeps to resume walks from. Each holds a copy of the catalog's maps of +/// the objects that the file changed before it (the objects themselves are shared), so they are +/// few. A request records two: at its statement, and at the statement of the last edit, which +/// covers typing in a statement with requests after it. +const WALK_CHECKPOINTS_MAX: usize = 4; + +/// A statement with the references that the walk resolved in it. +pub(crate) struct ResolvedStatement { + pub(crate) stmt: AnalysableStatement, + pub(crate) references: Vec, +} + +/// How an edit changed the document, as far as the walk is concerned. +#[derive(Debug, Clone, Copy)] +pub(crate) struct ContentChange { + /// The length of the text that is the same before and after the edit. + pub(crate) unchanged: TextSize, + /// The start of the first top-level statement, before the edit, that isn't the same after + /// it, or that doesn't end before the change. All statements before it are unchanged. + pub(crate) first_changed_statement: TextSize, +} + +/// The migration sections of a document, and the queries that run in an implicit transaction +/// block, as boundaries in file order. A walk to a statement depends on the boundaries before +/// it. Text after the statement can change them: a `StatementEnd` after it puts it in a query, +/// a statement added to its query makes the query run in an implicit block, and goose's +/// `NO TRANSACTION` applies to the whole file. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct MigrationContext { + /// Ordered by offset. + boundaries: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum Boundary { + Section { + start: TextSize, + down: bool, + transaction: bool, + }, + ImplicitBlockStart(TextSize), + ImplicitBlockEnd(TextSize), +} + +impl Boundary { + fn offset(&self) -> TextSize { + match *self { + Boundary::Section { start, .. } => start, + Boundary::ImplicitBlockStart(offset) | Boundary::ImplicitBlockEnd(offset) => offset, + } + } +} + +impl MigrationContext { + /// The context of `sections`. `statement_ranges` are the ranges of the top-level statements + /// of the document. + /// + /// Cost: O(log statements) per query of the sections. + pub(crate) fn new(sections: &[MigrationSection], statement_ranges: &[TextRange]) -> Self { + let mut boundaries: Vec = sections + .iter() + .flat_map(|section| { + std::iter::once(Boundary::Section { + start: section.range.start(), + down: section.direction == Direction::Down, + transaction: section.transaction, + }) + .chain(section.implicit_blocks(statement_ranges).flat_map( + |block| { + [ + Boundary::ImplicitBlockStart(block.start()), + Boundary::ImplicitBlockEnd(block.end()), + ] + }, + )) + }) + .collect(); + boundaries.sort_unstable_by_key(|boundary| (boundary.offset(), *boundary)); + Self { boundaries } + } + + /// The first offset from which the two contexts can differ: they have the same boundaries + /// before it. `None` if they are the same. + fn first_difference(&self, other: &Self) -> Option { + let same = self + .boundaries + .iter() + .zip(&other.boundaries) + .take_while(|(this, other)| this == other) + .count(); + self.boundaries + .get(same) + .into_iter() + .chain(other.boundaries.get(same)) + .map(Boundary::offset) + .min() + } +} + +#[derive(Default)] +pub(crate) struct WalkCache { + /// Least recently used first. + checkpoints: Vec, + /// The statements of each top-level statement that the walk resolved, by its start: the + /// top-level statement and those of its function body. Empty for a statement that doesn't + /// parse, or that has identifier parameters. + resolved: BTreeMap>, + /// The database snapshot as the statements of one top-level statement see it. + snapshot: Option, + /// The context of the document that the entries were computed in. + context: MigrationContext, +} + +struct SnapshotBefore { + /// The start of the top-level statement. + at: TextSize, + /// Whether it is the snapshot of the statements of its function body, which also see the + /// function itself. + in_body: bool, + snapshot: Arc, +} + +impl WalkCache { + /// The last checkpoint at or before `at` that fits `sections`. + pub(crate) fn checkpoint_before( + &mut self, + at: TextSize, + sections: &[MigrationSection], + ) -> Option { + let (index, _) = self + .checkpoints + .iter() + .enumerate() + .filter(|(_, checkpoint)| checkpoint.at() <= at && checkpoint.fits(sections)) + .max_by_key(|(_, checkpoint)| checkpoint.at())?; + // Most recently used last. + let checkpoint = self.checkpoints.remove(index); + self.checkpoints.push(checkpoint.clone()); + Some(checkpoint) + } + + /// Keeps `state`, replacing a checkpoint at the same statement, and drops the least + /// recently used one if there are too many. + pub(crate) fn record_checkpoint(&mut self, state: WalkState) { + self.checkpoints + .retain(|checkpoint| checkpoint.at() != state.at()); + if self.checkpoints.len() == WALK_CHECKPOINTS_MAX { + self.checkpoints.remove(0); + } + self.checkpoints.push(state); + } + + /// The resolved statements of the top-level statement that starts at `at`. + pub(crate) fn resolved(&self, at: TextSize) -> Option> { + self.resolved.get(&at).cloned() + } + + pub(crate) fn record_resolved(&mut self, at: TextSize, statements: Arc<[ResolvedStatement]>) { + self.resolved.insert(at, statements); + } + + /// The snapshot as the statements of the top-level statement at `at` see it, or the + /// statements of its function body if `in_body`. + pub(crate) fn snapshot(&self, at: TextSize, in_body: bool) -> Option> { + self.snapshot + .as_ref() + .filter(|cached| cached.at == at && cached.in_body == in_body) + .map(|cached| Arc::clone(&cached.snapshot)) + } + + pub(crate) fn record_snapshot(&mut self, at: TextSize, in_body: bool, snapshot: Arc) { + self.snapshot = Some(SnapshotBefore { + at, + in_body, + snapshot, + }); + } + + /// Takes out the entries that `context` invalidates, and returns them: those of the + /// statements from the first boundary where it differs from the context of the entries. + pub(crate) fn set_context(&mut self, context: &MigrationContext) -> WalkCache { + let Some(changed) = self.context.first_difference(context) else { + return WalkCache::default(); + }; + self.context = context.clone(); + let (checkpoints, invalid) = std::mem::take(&mut self.checkpoints) + .into_iter() + .partition(|checkpoint| checkpoint.at() < changed); + self.checkpoints = checkpoints; + WalkCache { + checkpoints: invalid, + resolved: self.resolved.split_off(&changed), + snapshot: self.snapshot.take_if(|cached| cached.at >= changed), + context: MigrationContext::default(), + } + } + + /// Takes out what `change` invalidates, and returns it. A checkpoint, and the snapshot of a + /// top-level statement, stay valid while the text before their statement and the statements + /// before it are the same. The resolved statements, and the snapshot of a function body, + /// also need their statement to be the same. + pub(crate) fn content_changed(&mut self, change: ContentChange) -> WalkCache { + let before_statement = + |at: TextSize| at <= change.unchanged && at <= change.first_changed_statement; + let (checkpoints, invalid) = std::mem::take(&mut self.checkpoints) + .into_iter() + .partition(|checkpoint| before_statement(checkpoint.at())); + self.checkpoints = checkpoints; + let resolved = self.resolved.split_off(&change.first_changed_statement); + let snapshot = self.snapshot.take_if(|cached| { + if cached.in_body { + cached.at >= change.first_changed_statement + } else { + !before_statement(cached.at) + } + }); + WalkCache { + checkpoints: invalid, + resolved, + snapshot, + context: MigrationContext::default(), + } + } + + pub(crate) fn is_empty(&self) -> bool { + self.checkpoints.is_empty() && self.resolved.is_empty() && self.snapshot.is_none() + } + + /// The number of resolved statements and checkpoints. + #[cfg(test)] + pub(crate) fn len(&self) -> (usize, usize) { + ( + self.resolved + .values() + .map(|statements| statements.len()) + .sum(), + self.checkpoints.len(), + ) + } +} diff --git a/crates/pgls_workspace/src/workspace/server/walk.rs b/crates/pgls_workspace/src/workspace/server/walk.rs new file mode 100644 index 000000000..bf82fdf9f --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/walk.rs @@ -0,0 +1,386 @@ +//! The walk over the statements of a document for the editor features: hover, completions, +//! inlay hints, semantic tokens and the expand-`*` action. +//! +//! It is lint's walk, [`pgls_analyser::walk_statements`], resumed from the checkpoints that the +//! document keeps, and what it resolves is kept for the next requests, see +//! [`super::document_cache::WalkCache`]. A request that only needs what is cached doesn't walk. + +use std::sync::Arc; + +use pgls_analyser::{AnalysableStatement, AnalysedFileContext, MigrationSection, WalkParams}; +use pgls_catalog::{ + Catalog, CatalogBase, Origin, Snapshot, + resolve::{ColumnSource, Reference, ResolveParams, resolve}, +}; +use pgls_text_size::{TextRange, TextSize}; + +use super::{ + WorkspaceServer, + document::{AnalyserDiagnosticsMapper, Document, DocumentStatement}, + document_cache::{CacheKey, MigrationContext, ResolvedStatement, WalkCache}, + document_sections, +}; +use crate::settings::Settings; + +/// What an editor request on a document walks with. +pub(super) struct EditorRequest<'a> { + /// The settings of the document. + pub(super) settings: Option<&'a Settings>, + /// The database snapshot the catalog starts from. + pub(super) snapshot: Arc, + pub(super) key: CacheKey, + /// The migration sections of the document. + sections: Vec, + context: MigrationContext, +} + +impl EditorRequest<'_> { + /// Runs `f` on what the walks over `doc` left for this request, see + /// [`super::document_cache::DocumentCache::with_walk`]. + fn with_walk(&self, doc: &Document, f: impl FnOnce(&mut WalkCache) -> R) -> R { + doc.cache().with_walk(&self.key, &self.context, f) + } +} + +impl WorkspaceServer { + /// What a request of `project_key` on `doc` with `settings` walks with, starting from the + /// database `snapshot`. + /// + /// Cost: O(lines) to find the migration sections of the document. + pub(super) fn editor_request<'a>( + &self, + project_key: crate::workspace::ProjectKey, + settings: Option<&'a Settings>, + snapshot: Arc, + doc: &Document, + ) -> EditorRequest<'a> { + let sections = settings + .map(|settings| document_sections(settings, doc)) + .unwrap_or_default(); + EditorRequest { + settings, + key: self.cache_key(project_key, Some(&snapshot)), + snapshot, + context: MigrationContext::new(§ions, doc.statement_ranges()), + sections, + } + } + + /// The indexed database snapshot, built once per snapshot. + pub(super) fn catalog_base(&self, snapshot: &Arc) -> Arc { + let mut cached = self.catalog_base.lock().unwrap(); + if let Some(base) = cached.as_ref().filter(|base| base.is_built_from(snapshot)) { + return Arc::clone(base); + } + let base = Arc::new(CatalogBase::new(Arc::clone(snapshot))); + *cached = Some(Arc::clone(&base)); + base + } + + /// Walks the statements of `doc` that start before `end`. The walk resumes from the last + /// checkpoint before the top-level statement that contains `start` or comes after it, and + /// records checkpoints there and at the statement where the last edit starts, so that the + /// next keystroke there resumes close to it. `visit` gets each statement with the file + /// context as it is before it. Returns the catalog as a statement at `end` sees it, and the + /// statements that were walked. + /// + /// Cost: O(statements between the checkpoint and `end`), plus a copy of the catalog's maps + /// for each checkpoint recorded. + fn walk( + &self, + request: &EditorRequest, + doc: &Document, + start: TextSize, + end: TextSize, + visit: impl FnMut(&AnalysableStatement, &AnalysedFileContext), + ) -> (Catalog, Vec) { + let search_path = request + .settings + .map(|settings| settings.typecheck.search_path.clone()) + .unwrap_or_default(); + let sections = &request.sections; + + let ranges = doc.statement_ranges(); + let request_at = ranges + .get(ranges.partition_point(|range| range.end() <= start)) + .map_or(end, |range| range.start()) + .min(end); + let edit_at = doc + .edited_at() + .and_then(|edited_at| { + let after = ranges.partition_point(|range| range.start() <= edited_at); + Some(ranges.get(after.checked_sub(1)?)?.start()) + }) + .filter(|edit_at| *edit_at < request_at); + let checkpoints: Vec = edit_at.into_iter().chain([request_at]).collect(); + let from = request.with_walk(doc, |walk| walk.checkpoint_before(request_at, sections)); + + let mut stmts = Vec::new(); + let mut unparsable = Vec::new(); + let walk_start = from.as_ref().map_or(TextSize::from(0), |from| from.at()); + for (stmt, _, skipped) in doc.iter_from(AnalyserDiagnosticsMapper, walk_start) { + let Some(stmt_start) = stmt + .as_ref() + .map(|stmt| stmt.range.start()) + .or_else(|| skipped.as_ref().map(|skipped| skipped.range.start())) + else { + continue; + }; + if stmt_start >= end { + break; + } + stmts.extend(stmt); + unparsable.extend(skipped); + } + #[cfg(test)] + self.walked_statements.fetch_add( + stmts.len() + unparsable.len(), + std::sync::atomic::Ordering::Relaxed, + ); + + let walk = pgls_analyser::walk_statements( + WalkParams { + stmts: &stmts, + unparsable: &unparsable, + statement_ranges: ranges, + catalog_base: Some(self.catalog_base(&request.snapshot)), + search_path: &request.snapshot.expand_search_path(&search_path), + sections, + end, + // Editor features show what names most likely refer to, so a `DO` block or a + // statement that doesn't parse doesn't turn them off for the rest of the file. + best_effort: true, + from, + checkpoints: &checkpoints, + }, + visit, + ); + request.with_walk(doc, |cache| { + for state in walk.checkpoints { + cache.record_checkpoint(state); + } + }); + (walk.catalog, stmts) + } + + /// The statements of the top-level statements that intersect `range` and start before its + /// end (or at it, for an empty range), with their references, top-level statement first. Statements that don't parse and statements with + /// identifier parameters, whose names are placeholders, have none. + /// + /// Cost: what is cached is looked up, O(log statements) each. Otherwise a walk from the + /// last checkpoint before the first statement that isn't cached to the end of `range`, see + /// [`Self::walk`]. + pub(super) fn resolved_statements( + &self, + request: &EditorRequest, + doc: &Document, + range: TextRange, + ) -> Vec> { + let ranges = doc.statement_ranges(); + let first = ranges.partition_point(|statement| statement.end() < range.start()); + // Like a walk to the end of `range`, which visits the statements that start before it, or + // at it if `range` is empty. + let last = ranges.partition_point(|statement| { + statement.start() < range.end() || statement.start() <= range.start() + }); + let tops = ranges.get(first..last).unwrap_or_default(); + let mut cached: Vec>> = request.with_walk(doc, |walk| { + tops.iter().map(|top| walk.resolved(top.start())).collect() + }); + let Some(first_missing) = cached.iter().position(Option::is_none) else { + return cached.into_iter().flatten().collect(); + }; + + let missing_start = tops[first_missing].start(); + let end = tops.last().map_or(range.end(), |top| top.end()); + // The index in `tops` of the top-level statement of a statement. + let top_of = |stmt: &AnalysableStatement| { + tops.partition_point(|top| top.start() <= stmt.range.start()) + .checked_sub(1) + }; + let mut references = Vec::new(); + let (_, stmts) = self.walk(request, doc, missing_start, end, |stmt, file_context| { + let Some(sql) = stmt.sql.as_deref() else { + return; + }; + if stmt.has_identifier_parameters + || top_of(stmt).is_none_or(|top| cached[top].is_some()) + { + return; + } + references.push(( + stmt.range, + resolve(ResolveParams { + stmt: &stmt.root, + catalog: file_context.catalog(), + search_path: file_context.session().search_path(), + function: stmt.function.as_ref(), + sql: Some(sql), + collect_references: true, + }) + .references, + )); + }); + + // The statements were visited in order, so the references are in the same order. + let mut resolved: Vec> = tops.iter().map(|_| Vec::new()).collect(); + let mut references = references.into_iter().peekable(); + for stmt in stmts { + let Some((_, stmt_references)) = + references.next_if(|(stmt_range, _)| *stmt_range == stmt.range) + else { + continue; + }; + if let Some(top) = top_of(&stmt) { + resolved[top].push(ResolvedStatement { + stmt, + references: stmt_references, + }); + } + } + request.with_walk(doc, |walk| { + for ((top, statements), cached) in tops.iter().zip(resolved).zip(&mut cached) { + if cached.is_none() { + let statements: Arc<[ResolvedStatement]> = statements.into(); + walk.record_resolved(top.start(), Arc::clone(&statements)); + *cached = Some(statements); + } + } + }); + cached.into_iter().flatten().collect() + } + + /// The database snapshot as `stmt` sees it: with the objects that the statements before it + /// create, change, or drop. The statements of a function body also see the function. + /// + /// Cost: a lookup if it was the last one built for the document. Otherwise a walk to the + /// statement, see [`Self::walk`], and building the snapshot, which is linear in the objects + /// that the file changed. + pub(super) fn snapshot_before( + &self, + request: &EditorRequest, + doc: &Document, + stmt: &DocumentStatement, + ) -> Arc { + let ranges = doc.statement_ranges(); + let top = ranges + .get(ranges.partition_point(|range| range.end() <= stmt.range.start())) + .map_or(stmt.range.start(), |range| range.start()); + let in_body = stmt.id.is_child(); + if let Some(snapshot) = request.with_walk(doc, |walk| walk.snapshot(top, in_body)) { + return snapshot; + } + let (catalog, _) = self.walk(request, doc, top, stmt.range.start(), |_, _| {}); + let snapshot = catalog + .snapshot() + .unwrap_or_else(|| Arc::clone(&request.snapshot)); + request.with_walk(doc, |walk| { + walk.record_snapshot(top, in_body, Arc::clone(&snapshot)) + }); + snapshot + } + + /// For hover on `stmt`: the references whose span contains `position`, and the database + /// snapshot as the statement sees it. + pub(super) fn hover_references( + &self, + request: &EditorRequest, + doc: &Document, + stmt: &DocumentStatement, + position: TextSize, + ) -> (Vec, Arc) { + let resolved = self.resolved_statements(request, doc, stmt.range); + let Some(resolved) = resolved + .iter() + .flat_map(|statements| statements.iter()) + .find(|resolved| resolved.stmt.range == stmt.range) + else { + // A statement that doesn't parse, or that has identifier parameters. + return (Vec::new(), self.snapshot_before(request, doc, stmt)); + }; + let position = resolved.stmt.statement_offset(position); + let (references, others): (Vec<_>, Vec<_>) = + resolved.references.iter().cloned().partition(|reference| { + reference + .span() + .is_some_and(|span| span.contains_inclusive(position)) + }); + // Objects that the file didn't change are the same in the database snapshot. + let snapshot = if in_database(&references, &others) { + Arc::clone(&request.snapshot) + } else { + self.snapshot_before(request, doc, stmt) + }; + (references, snapshot) + } + + /// The references whose span contains `position`, in the innermost statement at `position` + /// (the body of a SQL function rather than its `CREATE FUNCTION`). Their spans are relative + /// to the start of that statement, which is returned with them. Empty if the statement's + /// text isn't the source text, like a function body with escaped quotes, since the spans + /// would then not point into the document. + pub(super) fn references_at( + &self, + request: &EditorRequest, + doc: &Document, + position: TextSize, + ) -> (TextSize, Vec) { + let resolved = self.resolved_statements(request, doc, TextRange::empty(position)); + let Some(innermost) = resolved + .iter() + .flat_map(|statements| statements.iter()) + .rfind(|resolved| resolved.stmt.range.contains(position)) + .filter(|resolved| resolved.stmt.offsets.is_verbatim()) + else { + return (TextSize::default(), Vec::new()); + }; + let start = innermost.stmt.range.start(); + let position = position - start; + let references = innermost + .references + .iter() + .filter(|reference| { + reference + .span() + .is_some_and(|span| span.contains_inclusive(position)) + }) + .cloned() + .collect(); + (start, references) + } +} + +/// Whether hover finds what `references` resolve to in the database snapshot: objects that the +/// file didn't change, and columns of CTEs and subqueries. `others` are the other references of +/// the statement. +fn in_database(references: &[Reference], others: &[Reference]) -> bool { + // The statement references each relation that it takes columns from. + let unchanged_relation = |schema: &str, name: &str| { + references.iter().chain(others).any(|reference| { + matches!( + reference, + Reference::Relation { relation, .. } + if relation.schema == schema + && relation.name == name + && relation.origin == Origin::Database + ) + }) + }; + !references.is_empty() + && references.iter().all(|reference| match reference { + Reference::Relation { relation, .. } => relation.origin == Origin::Database, + Reference::Function { function, .. } => function.origin == Origin::Database, + Reference::Column { + source: ColumnSource::Relation { schema, name }, + .. + } => unchanged_relation(schema, name), + Reference::Column { .. } => true, + Reference::Type { type_info, .. } | Reference::EnumLiteral { type_info, .. } => { + type_info.origin == Origin::Database + } + // Parameters belong to a function of the file. + Reference::Parameter { .. } => false, + // The columns of a `*` can come from objects that the file changed. + Reference::Star { .. } => false, + }) +} diff --git a/docs/features/editor_features.md b/docs/features/editor_features.md index 4f52d639e..19ace2b46 100644 --- a/docs/features/editor_features.md +++ b/docs/features/editor_features.md @@ -108,4 +108,4 @@ The language server asks the editor to request inlay hints and semantic tokens a - The effects of `DO` blocks, `CALL`, `EXECUTE`, and statements that don't parse on the schema are unknown. The statements after them are resolved as if they changed nothing, so an object a `DO` block creates is not known, and one it drops still is. - Statements of SQL function bodies are resolved in `$$`, `'...'` and `E'...'` bodies, `BEGIN ATOMIC ... END` and `RETURN` bodies. Bodies in `U&'...'` strings and PL/pgSQL bodies are not covered. - No tokens for roles, for schemas outside qualified names, for DDL (column types in `create table`, function signatures), for CTE names, or for table aliases. Tokens that span several lines are skipped. -- Each request resolves the file up to the visible range. On very long files, like schema dumps with thousands of statements, hints and tokens take longer to appear. +- The first request on a file resolves it up to the visible range, and the next requests reuse what it resolved. An edit only invalidates what comes after it, so on very long files, like schema dumps with thousands of statements, hints and tokens take longer to appear after an edit near the start of the file.