nice, comes out as expected

Author: juleswritescode

crates/pg_completions/src/builder.rs

@@ -1,29 +1,30 @@
 use crate::{item::CompletionItem, CompletionResult};
 
-pub struct CompletionBuilder {
-    pub items: Vec<CompletionItem>,
+pub(crate) struct CompletionBuilder {
+    items: Vec<CompletionItem>,
 }
 
-pub struct CompletionConfig {}
+impl CompletionBuilder {
+    pub fn new() -> Self {
+        CompletionBuilder { items: vec![] }
+    }
 
-impl From<&CompletionConfig> for CompletionBuilder {
-    fn from(_config: &CompletionConfig) -> Self {
-        Self { items: Vec::new() }
+    pub fn add_item(&mut self, item: CompletionItem) {
+        self.items.push(item)
     }
-}
 
-impl CompletionBuilder {
     pub fn finish(mut self) -> CompletionResult {
         self.items.sort_by(|a, b| {
-            b.preselect
-                .cmp(&a.preselect)
-                .then_with(|| b.score.cmp(&a.score))
-                .then_with(|| a.data.label().cmp(b.data.label()))
+            b.score()
+                .cmp(&a.score())
+                .then_with(|| a.label.cmp(&b.label))
         });
 
-        self.items.dedup_by(|a, b| a.data.label() == b.data.label());
+        self.items.dedup_by(|a, b| a.label == b.label);
         self.items.truncate(crate::LIMIT);
+
         let Self { items, .. } = self;
+
         CompletionResult { items }
     }
 }

crates/pg_completions/src/complete.rs

@@ -1,6 +1,8 @@
 use text_size::TextSize;
 
-use crate::{builder, context::CompletionContext, item::CompletionItem};
+use crate::{
+    builder::CompletionBuilder, context::CompletionContext, item::CompletionItem, providers,
+};
 
 pub const LIMIT: usize = 50;
 
@@ -17,9 +19,18 @@ pub struct CompletionResult {
     pub items: Vec<CompletionItem>,
 }
 
-pub fn complete(params: &CompletionParams) -> CompletionResult {
-    let ctx = CompletionContext::new(params);
-    let mut builder = builder::CompletionBuilder::from(&builder::CompletionConfig {});
+pub fn complete(params: CompletionParams) -> CompletionResult {
+    let ctx = CompletionContext::new(&params);
+    let mut builder = CompletionBuilder::new();
+
+    if let Some(node) = ctx.ts_node {
+        match node.kind() {
+            "relation" => providers::complete_tables(&ctx, &mut builder),
+            _ => {}
+        }
+    } else {
+        // if query emtpy, autocomplete select keywords etc?
+    }
 
     builder.finish()
 }
@@ -55,7 +66,7 @@ mod tests {
             tree: Some(&tree),
         };
 
-        let result = complete(&p);
+        let result = complete(p);
 
         assert!(result.items.len() > 0);
     }
@@ -81,7 +92,7 @@ mod tests {
             tree: Some(&tree),
         };
 
-        let result = complete(&p);
+        let result = complete(p);
 
         assert!(result.items.len() > 0);
     }
@@ -120,7 +131,7 @@ mod tests {
             tree: Some(&tree),
         };
 
-        let result = complete(&p);
+        let result = complete(p);
 
         // TODO: actually assert that we get good autocompletion suggestions
         assert!(result.items.len() > 0);

crates/pg_completions/src/context.rs

@@ -1,34 +1,47 @@
 use pg_schema_cache::SchemaCache;
+use text_size::TextSize;
 
 use crate::CompletionParams;
 
-pub struct CompletionContext<'a> {
+pub(crate) struct CompletionContext<'a> {
     pub ts_node: Option<tree_sitter::Node<'a>>,
     pub tree: Option<&'a tree_sitter::Tree>,
     pub text: &'a str,
     pub schema_cache: &'a SchemaCache,
     pub original_token: Option<char>,
+    pub position: TextSize,
 }
 
 impl<'a> CompletionContext<'a> {
     pub fn new(params: &'a CompletionParams) -> Self {
+        let ts_node = find_ts_node(&params);
+
         Self {
-            ts_node: find_ts_node(params),
+            ts_node,
             tree: params.tree,
-            text: params.text,
+            text: &params.text,
             schema_cache: params.schema,
             original_token: find_original_token(params),
+            position: params.position,
+        }
+    }
+
+    pub fn get_ts_node_content(&self, ts_node: tree_sitter::Node<'a>) -> Option<&'a str> {
+        let source = self.text;
+        match ts_node.utf8_text(source.as_bytes()) {
+            Ok(content) => Some(content),
+            Err(_) => None,
         }
     }
 }
 
-fn find_original_token<'a>(params: &'a CompletionParams) -> Option<char> {
+fn find_original_token<'a>(params: &CompletionParams) -> Option<char> {
     let idx = usize::from(params.position);
     params.text.chars().nth(idx)
 }
 
 fn find_ts_node<'a>(params: &'a CompletionParams) -> Option<tree_sitter::Node<'a>> {
-    let tree = params.tree?;
+    let tree = params.tree.as_ref()?;
 
     let mut node = tree.root_node().named_descendant_for_byte_range(
         usize::from(params.position),

crates/pg_completions/src/item.rs

@@ -1,33 +1,41 @@
 use text_size::TextRange;
 
-#[derive(Debug, PartialEq, Eq)]
-pub enum CompletionItemData {
-    Table(pg_schema_cache::Table),
+use crate::relevance::CompletionRelevance;
+
+#[derive(Debug)]
+pub enum CompletionItemData<'a> {
+    Table(&'a pg_schema_cache::Table),
 }
 
-impl CompletionItemData {
-    pub fn label(&self) -> &str {
+impl<'a> CompletionItemData<'a> {
+    pub fn label(&self) -> String {
         match self {
-            CompletionItemData::Table(t) => t.name.as_str(),
+            CompletionItemData::Table(t) => t.name.clone(),
         }
     }
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug)]
 pub struct CompletionItem {
-    pub score: i32,
     pub range: TextRange,
-    pub preselect: bool,
-    pub data: CompletionItemData,
+    pub label: String,
+    relevance: CompletionRelevance,
 }
 
 impl CompletionItem {
-    pub fn new_simple(score: i32, range: TextRange, data: CompletionItemData) -> Self {
+    pub(crate) fn new(
+        range: TextRange,
+        data: CompletionItemData,
+        relevance: CompletionRelevance,
+    ) -> Self {
         Self {
-            score,
             range,
-            preselect: false,
-            data,
+            label: data.label(),
+            relevance,
         }
     }
+
+    pub(crate) fn score(&self) -> i32 {
+        self.relevance.score()
+    }
 }

crates/pg_completions/src/lib.rs

@@ -2,6 +2,7 @@ mod builder;
 mod complete;
 mod context;
 mod item;
+mod providers;
 mod relevance;
 
 pub use complete::*;

crates/pg_completions/src/providers/mod.rs

@@ -0,0 +1,3 @@
+mod tables;
+
+pub use tables::*;

crates/pg_completions/src/providers/tables.rs

@@ -1,31 +1,38 @@
+use pg_schema_cache::Table;
 use text_size::{TextRange, TextSize};
 
-use crate::{builder::CompletionBuilder, CompletionItem, CompletionItemData};
-
-use super::CompletionProviderParams;
-
-// todo unify this in a type resolver crate
-pub fn complete_tables<'a>(
-    params: CompletionProviderParams<'a>,
-    builder: &mut CompletionBuilder<'a>,
-) {
-    if let Some(ts) = params.ts_node {
-        let range = TextRange::new(
-            TextSize::try_from(ts.start_byte()).unwrap(),
-            TextSize::try_from(ts.end_byte()).unwrap(),
-        );
-        match ts.kind() {
-            "relation" => {
-                // todo better search
-                params.schema.tables.iter().for_each(|table| {
-                    builder.items.push(CompletionItem::new_simple(
-                        1,
-                        range,
-                        CompletionItemData::Table(table),
-                    ));
-                });
-            }
-            _ => {}
-        }
+use crate::{
+    builder::CompletionBuilder,
+    context::CompletionContext,
+    item::{CompletionItem, CompletionItemData},
+    relevance::CompletionRelevance,
+};
+
+pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+    let available_tables = &ctx.schema_cache.tables;
+
+    let completion_items: Vec<CompletionItem> = available_tables
+        .iter()
+        .map(|table| to_completion_item(ctx, table))
+        .collect();
+
+    for item in completion_items {
+        builder.add_item(item);
     }
 }
+
+fn to_completion_item(ctx: &CompletionContext, table: &Table) -> CompletionItem {
+    let data = CompletionItemData::Table(table);
+
+    let start = ctx.position;
+    let end = start + TextSize::from(table.name.len() as u32);
+    let range = TextRange::new(start, end);
+
+    let mut relevance = CompletionRelevance::default();
+
+    relevance.set_is_catalog(&table.schema);
+    relevance.set_matches_prefix(ctx, &table.name);
+    relevance.set_matches_schema(ctx, &table.schema);
+
+    CompletionItem::new(range, data, relevance)
+}

crates/pg_completions/src/relevance.rs

@@ -1 +1,54 @@
-struct CompletionItemScore {}
+use crate::context::CompletionContext;
+
+#[derive(Debug, Default)]
+pub(crate) struct CompletionRelevance {
+    /// does the underlying data match the expected schema we can determine from the query?
+    matches_schema: bool,
+
+    /// Is the underlying item from the pg_catalog schema?
+    is_catalog: bool,
+
+    /// Do the characters the users typed match at least the first 3 characters
+    /// of the underlying data's name?
+    matches_prefix: usize,
+}
+
+impl CompletionRelevance {
+    pub fn score(&self) -> i32 {
+        let mut score: i32 = 0;
+
+        if self.matches_schema {
+            score += 5;
+        } else if self.is_catalog {
+            score -= 1;
+        }
+
+        score += (self.matches_prefix * 5) as i32;
+
+        score
+    }
+
+    pub fn set_matches_schema(&mut self, ctx: &CompletionContext, schema: &str) {
+        let node = ctx.ts_node.unwrap();
+        self.matches_schema = node
+            .prev_named_sibling()
+            .is_some_and(|n| ctx.get_ts_node_content(n).is_some_and(|c| c == schema));
+    }
+
+    pub fn set_is_catalog(&mut self, schema: &str) {
+        self.is_catalog = schema == "pg_catalog"
+    }
+
+    pub fn set_matches_prefix(&mut self, ctx: &CompletionContext, name: &str) {
+        let node = ctx.ts_node.unwrap();
+
+        let content = match ctx.get_ts_node_content(node) {
+            Some(c) => c,
+            None => return,
+        };
+
+        if name.starts_with(content) {
+            self.matches_prefix = content.len();
+        };
+    }
+}

crates/pg_lsp/src/session.rs

@@ -235,17 +235,21 @@ impl Session {
 
         let schema_cache = ide.schema_cache.read().expect("No Schema Cache");
 
-        let completion_items = pg_completions::complete(&CompletionParams {
+        let completion_items = pg_completions::complete(CompletionParams {
             position: offset - range.start() - TextSize::from(1),
-            text: stmt.text.as_str(),
-            tree: ide.tree_sitter.tree(&stmt).as_ref().map(|x| x.as_ref()),
+            text: &stmt.text,
+            tree: ide
+                .tree_sitter
+                .tree(&stmt)
+                .as_ref()
+                .and_then(|t| Some(t.as_ref())),
             schema: &schema_cache,
         })
         .items
         .into_iter()
         .map(|i| CompletionItem {
             // TODO: add more data
-            label: i.data.label().to_string(),
+            label: i.label,
             label_details: None,
             kind: Some(CompletionItemKind::CLASS),
             detail: None,