perf(workspace): find the statement for completions without going over the whole file

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<tree_sitter::Tree>)> {
-    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:?}"
+            );
+        }
+    }
 }

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