From 894654d325170a4382d21342c63e9dcbdf578e01 Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Sun, 25 Oct 2026 18:43:59 +0000
Subject: [PATCH] perf(workspace): find the statement for completions without
 going over the whole file

get_statement_for_completions counted every statement of the document, which maps each of
them and looks up its parse result, and then filtered all statements for the ones around the
cursor. It only needs to know whether there are none, one or more statements, and the filter
only accepts statements that contain the cursor or end up to two characters before it, so it
now looks at those only.
---
 .../src/features/completions.rs               | 53 +++++++++++++++----
 .../src/workspace/server/document.rs          | 27 +++++++++-
 2 files changed, 68 insertions(+), 12 deletions(-)

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<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:?}"
+            );
+        }
+    }
 }
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,
+                },
+            )
+        }
     }
 }