ide: inlay hints for insert & add goto def for hints (#771)

Author: sbdchd

crates/squawk_ide/src/inlay_hints.rs

@@ -1,7 +1,8 @@
 use crate::binder;
 use crate::binder::Binder;
 use crate::resolve;
-use rowan::TextSize;
+use crate::symbols::Name;
+use rowan::{TextRange, TextSize};
 use squawk_syntax::ast::{self, AstNode};
 
 /// `VSCode` has some theming options based on these types.
@@ -16,15 +17,18 @@ pub struct InlayHint {
     pub position: TextSize,
     pub label: String,
     pub kind: InlayHintKind,
+    pub target: Option<TextRange>,
 }
 
 pub fn inlay_hints(file: &ast::SourceFile) -> Vec<InlayHint> {
     let mut hints = vec![];
     let binder = binder::bind(file);
 
     for node in file.syntax().descendants() {
-        if let Some(call_expr) = ast::CallExpr::cast(node) {
+        if let Some(call_expr) = ast::CallExpr::cast(node.clone()) {
             inlay_hint_call_expr(&mut hints, file, &binder, call_expr);
+        } else if let Some(insert) = ast::Insert::cast(node) {
+            inlay_hint_insert(&mut hints, file, &binder, insert);
         }
     }
 
@@ -59,10 +63,12 @@ fn inlay_hint_call_expr(
         for (param, arg) in param_list.params().zip(arg_list.args()) {
             if let Some(param_name) = param.name() {
                 let arg_start = arg.syntax().text_range().start();
+                let target = Some(param_name.syntax().text_range());
                 hints.push(InlayHint {
                     position: arg_start,
                     label: format!("{}: ", param_name.syntax().text()),
                     kind: InlayHintKind::Parameter,
+                    target,
                 });
             }
         }
@@ -71,6 +77,62 @@ fn inlay_hint_call_expr(
     Some(())
 }
 
+fn inlay_hint_insert(
+    hints: &mut Vec<InlayHint>,
+    file: &ast::SourceFile,
+    binder: &Binder,
+    insert: ast::Insert,
+) -> Option<()> {
+    let values = insert.values()?;
+    let row_list = values.row_list()?;
+
+    let columns: Vec<(Name, Option<TextRange>)> = if let Some(column_list) = insert.column_list() {
+        let table_arg_list = resolve::resolve_insert_table_columns(file, binder, &insert);
+
+        column_list
+            .columns()
+            .filter_map(|col| {
+                let col_name = resolve::extract_column_name(&col)?;
+                let target = table_arg_list
+                    .as_ref()
+                    .and_then(|list| resolve::find_column_in_table(list, &col_name));
+                Some((col_name, target))
+            })
+            .collect()
+    } else {
+        let table_arg_list = resolve::resolve_insert_table_columns(file, binder, &insert)?;
+
+        table_arg_list
+            .args()
+            .filter_map(|arg| {
+                if let ast::TableArg::Column(column) = arg
+                    && let Some(name) = column.name()
+                {
+                    let col_name = Name::new(name.syntax().text().to_string());
+                    let target = Some(name.syntax().text_range());
+                    Some((col_name, target))
+                } else {
+                    None
+                }
+            })
+            .collect()
+    };
+
+    for row in row_list.rows() {
+        for ((column_name, target), expr) in columns.iter().zip(row.exprs()) {
+            let expr_start = expr.syntax().text_range().start();
+            hints.push(InlayHint {
+                position: expr_start,
+                label: format!("{}: ", column_name),
+                kind: InlayHintKind::Parameter,
+                target: *target,
+            });
+        }
+    }
+
+    Some(())
+}
+
 #[cfg(test)]
 mod test {
     use crate::inlay_hints::inlay_hints;
@@ -216,4 +278,68 @@ select foo(1, 2);
           ╰╴           ───
         ");
     }
+
+    #[test]
+    fn insert_with_column_list() {
+        assert_snapshot!(check_inlay_hints("
+create table t (column_a int, column_b int, column_c text);
+insert into t (column_a, column_c) values (1, 'foo');
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ insert into t (column_a, column_c) values (column_a: 1, column_c: 'foo');
+          ╰╴                                           ──────────   ──────────
+        ");
+    }
+
+    #[test]
+    fn insert_without_column_list() {
+        assert_snapshot!(check_inlay_hints("
+create table t (column_a int, column_b int, column_c text);
+insert into t values (1, 2, 'foo');
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ insert into t values (column_a: 1, column_b: 2, column_c: 'foo');
+          ╰╴                      ──────────   ──────────   ──────────
+        ");
+    }
+
+    #[test]
+    fn insert_multiple_rows() {
+        assert_snapshot!(check_inlay_hints("
+create table t (x int, y int);
+insert into t values (1, 2), (3, 4);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ insert into t values (x: 1, y: 2), (x: 3, y: 4);
+          ╰╴                      ───   ───     ───   ───
+        ");
+    }
+
+    #[test]
+    fn insert_no_create_table() {
+        assert_snapshot!(check_inlay_hints("
+insert into t (a, b) values (1, 2);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        2 │ insert into t (a, b) values (a: 1, b: 2);
+          ╰╴                             ───   ───
+        ");
+    }
+
+    #[test]
+    fn insert_more_values_than_columns() {
+        assert_snapshot!(check_inlay_hints("
+create table t (a int, b int);
+insert into t values (1, 2, 3);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ insert into t values (a: 1, b: 2, 3);
+          ╰╴                      ───   ───
+        ");
+    }
 }

crates/squawk_ide/src/resolve.rs

@@ -1,4 +1,4 @@
-use rowan::TextSize;
+use rowan::{TextRange, TextSize};
 use squawk_syntax::{
     SyntaxNodePtr,
     ast::{self, AstNode},
@@ -318,6 +318,53 @@ fn extract_schema_name(path: &ast::Path) -> Option<Schema> {
         .map(|name_ref| Schema(Name::new(name_ref.syntax().text().to_string())))
 }
 
+pub(crate) fn extract_column_name(col: &ast::Column) -> Option<Name> {
+    let text = if let Some(name_ref) = col.name_ref() {
+        name_ref.syntax().text().to_string()
+    } else {
+        let name = col.name()?;
+        name.syntax().text().to_string()
+    };
+    Some(Name::new(text))
+}
+
+pub(crate) fn find_column_in_table(
+    table_arg_list: &ast::TableArgList,
+    col_name: &Name,
+) -> Option<TextRange> {
+    table_arg_list.args().find_map(|arg| {
+        if let ast::TableArg::Column(column) = arg
+            && let Some(name) = column.name()
+            && Name::new(name.syntax().text().to_string()) == *col_name
+        {
+            Some(name.syntax().text_range())
+        } else {
+            None
+        }
+    })
+}
+
+pub(crate) fn resolve_insert_table_columns(
+    file: &ast::SourceFile,
+    binder: &Binder,
+    insert: &ast::Insert,
+) -> Option<ast::TableArgList> {
+    let path = insert.path()?;
+    let table_name = extract_table_name(&path)?;
+    let schema = extract_schema_name(&path);
+    let position = insert.syntax().text_range().start();
+
+    let table_ptr = resolve_table(binder, &table_name, &schema, position)?;
+    let root = file.syntax();
+    let table_name_node = table_ptr.to_node(root);
+
+    let create_table = table_name_node
+        .ancestors()
+        .find_map(ast::CreateTable::cast)?;
+
+    create_table.table_arg_list()
+}
+
 pub(crate) fn resolve_table_info(binder: &Binder, path: &ast::Path) -> Option<(Schema, String)> {
     let table_name_str = extract_table_name_from_path(path)?;
     let schema = extract_schema_from_path(path);

crates/squawk_ide/src/symbols.rs

@@ -29,6 +29,12 @@ impl Name {
     }
 }
 
+impl fmt::Display for Name {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 fn normalize_identifier(text: &str) -> SmolStr {
     if text.starts_with('"') && text.ends_with('"') && text.len() >= 2 {
         text[1..text.len() - 1].into()

crates/squawk_server/src/lib.rs

@@ -8,9 +8,10 @@ use lsp_types::{
     DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
     GotoDefinitionParams, GotoDefinitionResponse, Hover, HoverContents, HoverParams,
     HoverProviderCapability, InitializeParams, InlayHint, InlayHintKind, InlayHintLabel,
-    InlayHintParams, LanguageString, Location, MarkedString, OneOf, PublishDiagnosticsParams,
-    ReferenceParams, SelectionRangeParams, SelectionRangeProviderCapability, ServerCapabilities,
-    TextDocumentSyncCapability, TextDocumentSyncKind, Url, WorkDoneProgressOptions, WorkspaceEdit,
+    InlayHintLabelPart, InlayHintParams, LanguageString, Location, MarkedString, OneOf,
+    PublishDiagnosticsParams, ReferenceParams, SelectionRangeParams,
+    SelectionRangeProviderCapability, ServerCapabilities, TextDocumentSyncCapability,
+    TextDocumentSyncKind, Url, WorkDoneProgressOptions, WorkspaceEdit,
     notification::{
         DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, Notification as _,
         PublishDiagnostics,
@@ -257,9 +258,24 @@ fn handle_inlay_hints(
                 squawk_ide::inlay_hints::InlayHintKind::Type => InlayHintKind::TYPE,
                 squawk_ide::inlay_hints::InlayHintKind::Parameter => InlayHintKind::PARAMETER,
             };
+
+            let label = if let Some(target_range) = hint.target {
+                InlayHintLabel::LabelParts(vec![InlayHintLabelPart {
+                    value: hint.label,
+                    location: Some(Location {
+                        uri: uri.clone(),
+                        range: lsp_utils::range(&line_index, target_range),
+                    }),
+                    tooltip: None,
+                    command: None,
+                }])
+            } else {
+                InlayHintLabel::String(hint.label)
+            };
+
             InlayHint {
                 position,
-                label: InlayHintLabel::String(hint.label),
+                label,
                 kind: Some(kind),
                 text_edits: None,
                 tooltip: None,