give it a spin

Author: juleswritescode

crates/pgt_workspace/src/features/completions.rs

@@ -1,6 +1,9 @@
+use itertools::Itertools;
 use pgt_completions::CompletionItem;
 use pgt_fs::PgTPath;
-use pgt_text_size::TextSize;
+use pgt_text_size::{TextRange, TextSize};
+
+use crate::workspace::{Document, Statement, StatementId};
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
@@ -24,3 +27,56 @@ impl IntoIterator for CompletionsResult {
         self.items.into_iter()
     }
 }
+
+pub(crate) fn get_statement_for_completions<'a>(
+    doc: &'a Document,
+    position: TextSize,
+) -> Option<(Statement, &'a TextRange, &'a str)> {
+    let count = doc.statement_count();
+    // no arms no cookies
+    if count == 0 {
+        return None;
+    }
+
+    /*
+     * We allow an offset of two for the statement:
+     *
+     * select * from | <-- we want to suggest items for the next token.
+     *
+     * However, if the current statement is terminated by a semicolon, we don't apply any
+     * offset.
+     *
+     * select * from users; | <-- no autocompletions here.
+     */
+    let matches_expanding_range = |stmt_id: StatementId, range: &TextRange, position: TextSize| {
+        let measuring_range = if doc.is_terminated_by_semicolon(stmt_id).unwrap() {
+            *range
+        } else {
+            range.checked_expand_end(2.into()).unwrap_or(*range)
+        };
+        measuring_range.contains(position)
+    };
+
+    if count == 1 {
+        let (stmt, range, txt) = doc.iter_statements_with_text_and_range().next().unwrap();
+        if matches_expanding_range(stmt.id, range, position) {
+            Some((stmt, range, txt))
+        } else {
+            None
+        }
+    } else {
+        /*
+         * If we have multiple statements, we want to make sure that we do not overlap
+         * with the next one.
+         *
+         * select 1 |select 1;
+         */
+        let mut stmts = doc.iter_statements_with_text_and_range().tuple_windows();
+        stmts
+            .find(|((current_stmt, rcurrent, _), (_, rnext, _))| {
+                let overlaps_next = rnext.contains(position);
+                matches_expanding_range(current_stmt.id, rcurrent, position) && !overlaps_next
+            })
+            .map(|t| t.0)
+    }
+}

crates/pgt_workspace/src/workspace.rs

@@ -21,7 +21,7 @@ use crate::{
 mod client;
 mod server;
 
-pub(crate) use server::StatementId;
+pub(crate) use server::document::{Document, Statement, StatementId};
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]

crates/pgt_workspace/src/workspace/server.rs

@@ -5,16 +5,13 @@ use async_helper::run_async;
 use change::StatementChange;
 use dashmap::DashMap;
 use db_connection::DbConnection;
-pub(crate) use document::StatementId;
 use document::{Document, Statement};
 use futures::{StreamExt, stream};
-use itertools::Itertools;
 use pg_query::PgQueryStore;
 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;
@@ -29,7 +26,7 @@ use crate::{
             self, CodeAction, CodeActionKind, CodeActionsResult, CommandAction,
             CommandActionCategory, ExecuteStatementParams, ExecuteStatementResult,
         },
-        completions::{CompletionsResult, GetCompletionsParams},
+        completions::{self, CompletionsResult, GetCompletionsParams},
         diagnostics::{PullDiagnosticsParams, PullDiagnosticsResult},
     },
     settings::{Settings, SettingsHandle, SettingsHandleMut},
@@ -44,7 +41,7 @@ mod analyser;
 mod async_helper;
 mod change;
 mod db_connection;
-mod document;
+pub(crate) mod document;
 mod migration;
 mod pg_query;
 mod schema_cache_manager;
@@ -540,64 +537,12 @@ impl Workspace for WorkspaceServer {
             .get(&params.path)
             .ok_or(WorkspaceError::not_found())?;
 
-        let count = doc.statement_count();
-        // no arms no cookies
-        if count == 0 {
-            return Ok(CompletionsResult::default());
-        }
-
-        /*
-         * We allow an offset of two for the statement:
-         *
-         * select * from | <-- we want to suggest items for the next token.
-         *
-         * However, if the current statement is terminated by a semicolon, we don't apply any
-         * offset.
-         *
-         * select * from users; | <-- no autocompletions here.
-         */
-        let matches_expanding_range =
-            |stmt_id: StatementId, range: &TextRange, position: TextSize| {
-                let measuring_range = if doc.is_terminated_by_semicolon(stmt_id).unwrap() {
-                    *range
-                } else {
-                    range.checked_expand_end(2.into()).unwrap_or(*range)
-                };
-                measuring_range.contains(position)
+        let (statement, stmt_range, text) =
+            match completions::get_statement_for_completions(&doc, params.position) {
+                None => return Ok(CompletionsResult::default()),
+                Some(s) => s,
             };
 
-        let maybe_statement = if count == 1 {
-            let (stmt, range, txt) = doc.iter_statements_with_text_and_range().next().unwrap();
-            if matches_expanding_range(stmt.id, range, params.position) {
-                Some((stmt, range, txt))
-            } else {
-                None
-            }
-        } else {
-            /*
-             * If we have multiple statements, we want to make sure that we do not overlap
-             * with the next one.
-             *
-             * select 1 |select 1;
-             */
-            let mut stmts = doc.iter_statements_with_text_and_range().tuple_windows();
-            stmts
-                .find(|((current_stmt, rcurrent, _), (_, rnext, _))| {
-                    let overlaps_next = rnext.contains(params.position);
-                    matches_expanding_range(current_stmt.id, &rcurrent, params.position)
-                        && !overlaps_next
-                })
-                .map(|t| t.0)
-        };
-
-        let (statement, stmt_range, text) = match maybe_statement {
-            Some(it) => it,
-            None => {
-                tracing::debug!("No matching statement found for completion.");
-                return Ok(CompletionsResult::default());
-            }
-        };
-
         // `offset` is the position in the document,
         // but we need the position within the *statement*.
         let position = params.position - stmt_range.start();