diff --git a/crates/pgls_workspace/src/features/completions.rs b/crates/pgls_workspace/src/features/completions.rs index 85a242091..c16833378 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::{TextRange, TextSize}; -use crate::workspace::{Document, GetCompletionsFilter, StatementId, WithCSTMapper}; +use crate::workspace::{DefaultMapper, Document, StatementId, WithCSTMapper}; #[derive(Debug, serde::Serialize, serde::Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -45,18 +45,14 @@ pub(crate) fn get_statement_for_completions( doc: &Document, position: TextSize, ) -> Option<(StatementId, TextRange, 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() @@ -90,7 +86,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; @@ -149,7 +145,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()); } @@ -228,4 +224,41 @@ 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 looked only at the cursor. + let mut expected = None; + for current in statements.by_ref() { + if expected.as_ref().is_some_and( + |(previous, _, _): &(crate::workspace::StatementId, _, _)| { + current.1.contains(position) && !current.0.is_child_of(previous) + }, + ) { + expected = None; + break; + } + expected = Some(current); + } + assert_eq!( + get_statement_for_completions(&doc, position).map(|(id, range, _)| (id, range)), + expected.map(|(id, range, _)| (id, range)), + "at {position:?}" + ); + } + } } diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 32d8069e5..93d77ae35 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -215,8 +215,31 @@ impl Document { } } - pub fn count(&self) -> usize { - self.iter(DefaultMapper).count() + /// 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.split.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, + }, + ) + } } }