hmmm

Author: juleswritescode

crates/pgt_completions/src/complete.rs

@@ -17,7 +17,10 @@ pub struct CompletionParams<'a> {
     pub tree: Option<&'a tree_sitter::Tree>,
 }
 
-#[tracing::instrument(level = "debug")]
+#[tracing::instrument(level = "debug", skip_all, fields(
+    text = params.text,
+    position = params.position.to_string()
+))]
 pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
     let ctx = CompletionContext::new(&params);
 

crates/pgt_completions/src/context.rs

@@ -49,12 +49,29 @@ impl TryFrom<String> for ClauseType {
 }
 
 pub(crate) struct CompletionContext<'a> {
-    pub ts_node: Option<tree_sitter::Node<'a>>,
+    pub node_under_cursor: Option<tree_sitter::Node<'a>>,
+    pub previous_node: Option<tree_sitter::Node<'a>>,
+
     pub tree: Option<&'a tree_sitter::Tree>,
     pub text: &'a str,
     pub schema_cache: &'a SchemaCache,
     pub position: usize,
 
+    /// If the cursor of the user is offset to the right of the statement,
+    /// we'll have to move it back to the last node, otherwise, tree-sitter will break.
+    /// However, knowing that the user is typing on the "next" node lets us prioritize different completion results.
+    /// We consider an offset of up to two characters as valid.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// select * from {}
+    /// ```
+    ///
+    /// We'll adjust the cursor position so it lies on the "from" token – but we're looking
+    /// for table completions.
+    pub cursor_offset_from_end: bool,
+
     pub schema_name: Option<String>,
     pub wrapping_clause_type: Option<ClauseType>,
     pub is_invocation: bool,
@@ -70,7 +87,9 @@ impl<'a> CompletionContext<'a> {
             text: &params.text,
             schema_cache: params.schema,
             position: usize::from(params.position),
-            ts_node: None,
+            cursor_offset_from_end: false,
+            previous_node: None,
+            node_under_cursor: None,
             schema_name: None,
             wrapping_clause_type: None,
             wrapping_statement_range: None,
@@ -81,8 +100,6 @@ impl<'a> CompletionContext<'a> {
         ctx.gather_tree_context();
         ctx.gather_info_from_ts_queries();
 
-        println!("Here's my node: {:?}", ctx.ts_node.unwrap());
-
         ctx
     }
 
@@ -147,30 +164,34 @@ impl<'a> CompletionContext<'a> {
          * `select * from use           {}` becomes `select * from use{}`.
          */
         let current_node = cursor.node();
+        let position_cache = self.position.clone();
         while cursor.goto_first_child_for_byte(self.position).is_none() && self.position > 0 {
             self.position -= 1;
         }
 
+        let cursor_offset = position_cache - self.position;
+        self.cursor_offset_from_end = cursor_offset > 0 && cursor_offset <= 2;
+
         self.gather_context_from_node(cursor, current_node);
     }
 
     fn gather_context_from_node(
         &mut self,
         mut cursor: tree_sitter::TreeCursor<'a>,
-        previous_node: tree_sitter::Node<'a>,
+        parent_node: tree_sitter::Node<'a>,
     ) {
         let current_node = cursor.node();
 
         // prevent infinite recursion – this can happen if we only have a PROGRAM node
-        if current_node.kind() == previous_node.kind() {
-            self.ts_node = Some(current_node);
+        if current_node.kind() == parent_node.kind() {
+            self.node_under_cursor = Some(current_node);
             return;
         }
 
-        match previous_node.kind() {
+        match parent_node.kind() {
             "statement" | "subquery" => {
                 self.wrapping_clause_type = current_node.kind().try_into().ok();
-                self.wrapping_statement_range = Some(previous_node.range());
+                self.wrapping_statement_range = Some(parent_node.range());
             }
             "invocation" => self.is_invocation = true,
 
@@ -202,7 +223,23 @@ impl<'a> CompletionContext<'a> {
 
         // We have arrived at the leaf node
         if current_node.child_count() == 0 {
-            self.ts_node = Some(current_node);
+            if self.cursor_offset_from_end {
+                self.node_under_cursor = None;
+                self.previous_node = Some(current_node);
+            } else {
+                // for the previous node, either select the previous sibling,
+                // or collect the parent's previous sibling's last child.
+                let previous = match current_node.prev_sibling() {
+                    Some(n) => Some(n),
+                    None => {
+                        let sib_of_parent = parent_node.prev_sibling();
+                        sib_of_parent.and_then(|p| p.children(&mut cursor).last())
+                    }
+                };
+                self.node_under_cursor = Some(current_node);
+                self.previous_node = previous;
+            }
+
             return;
         }
 
@@ -361,7 +398,7 @@ mod tests {
 
             let ctx = CompletionContext::new(&params);
 
-            let node = ctx.ts_node.unwrap();
+            let node = ctx.node_under_cursor.unwrap();
 
             assert_eq!(ctx.get_ts_node_content(node), Some("select"));
 
@@ -389,7 +426,7 @@ mod tests {
 
         let ctx = CompletionContext::new(&params);
 
-        let node = ctx.ts_node.unwrap();
+        let node = ctx.node_under_cursor.unwrap();
 
         assert_eq!(ctx.get_ts_node_content(node), Some("from"));
         assert_eq!(
@@ -415,7 +452,7 @@ mod tests {
 
         let ctx = CompletionContext::new(&params);
 
-        let node = ctx.ts_node.unwrap();
+        let node = ctx.node_under_cursor.unwrap();
 
         assert_eq!(ctx.get_ts_node_content(node), Some(""));
         assert_eq!(ctx.wrapping_clause_type, None);
@@ -440,7 +477,7 @@ mod tests {
 
         let ctx = CompletionContext::new(&params);
 
-        let node = ctx.ts_node.unwrap();
+        let node = ctx.node_under_cursor.unwrap();
 
         assert_eq!(ctx.get_ts_node_content(node), Some("fro"));
         assert_eq!(ctx.wrapping_clause_type, Some(ClauseType::Select));

crates/pgt_completions/src/relevance.rs

@@ -41,7 +41,10 @@ impl CompletionRelevance<'_> {
     }
 
     fn check_matches_query_input(&mut self, ctx: &CompletionContext) {
-        let node = ctx.ts_node.unwrap();
+        let node = match ctx.node_under_cursor {
+            Some(node) => node,
+            None => return,
+        };
 
         let content = match ctx.get_ts_node_content(node) {
             Some(c) => c,

crates/pgt_workspace/src/workspace/server.rs

@@ -13,6 +13,7 @@ use pgt_analyse::{AnalyserOptions, AnalysisFilter};
 use pgt_analyser::{Analyser, AnalyserConfig, AnalyserContext};
 use pgt_diagnostics::{Diagnostic, DiagnosticExt, Severity, serde::Diagnostic as SDiagnostic};
 use pgt_fs::{ConfigName, PgTPath};
+use pgt_text_size::{TextRange, TextSize};
 use pgt_typecheck::TypecheckParams;
 use schema_cache_manager::SchemaCacheManager;
 use sqlx::Executor;
@@ -535,13 +536,18 @@ impl Workspace for WorkspaceServer {
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
-        let (statement, stmt_range, text) = match doc
-            .iter_statements_with_text_and_range()
-            .find(|(_, r, _)| r.contains(params.position))
-        {
-            Some(s) => s,
-            None => return Ok(CompletionsResult::default()),
-        };
+        let (statement, stmt_range, text) =
+            match doc.iter_statements_with_text_and_range().find(|(_, r, _)| {
+                let expanded_range = TextRange::new(
+                    r.start(),
+                    r.end().checked_add(TextSize::new(2)).unwrap_or(r.end()),
+                );
+
+                expanded_range.contains(params.position)
+            }) {
+                Some(s) => s,
+                None => return Ok(CompletionsResult::default()),
+            };
 
         // `offset` is the position in the document,
         // but we need the position within the *statement*.