ide: inlay hints for function calls (#769)

Author: sbdchd

PLAN.md

@@ -568,6 +568,15 @@ select a, b from t
 group by a, b;
 ```
 
+### Rule: unresolved column
+
+```sql
+create function foo(a int, b int) returns int
+as 'select $0'
+--         ^^ unresolved column, did you mean `a` or `b`?
+language sql;
+```
+
 ### Rule: unused column
 
 ```sql

crates/squawk_ide/src/inlay_hints.rs

@@ -0,0 +1,219 @@
+use crate::binder;
+use crate::binder::Binder;
+use crate::resolve;
+use rowan::TextSize;
+use squawk_syntax::ast::{self, AstNode};
+
+/// `VSCode` has some theming options based on these types.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum InlayHintKind {
+    Type,
+    Parameter,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct InlayHint {
+    pub position: TextSize,
+    pub label: String,
+    pub kind: InlayHintKind,
+}
+
+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) {
+            inlay_hint_call_expr(&mut hints, file, &binder, call_expr);
+        }
+    }
+
+    hints
+}
+
+fn inlay_hint_call_expr(
+    hints: &mut Vec<InlayHint>,
+    file: &ast::SourceFile,
+    binder: &Binder,
+    call_expr: ast::CallExpr,
+) -> Option<()> {
+    let arg_list = call_expr.arg_list()?;
+    let expr = call_expr.expr()?;
+
+    let name_ref = if let Some(name_ref) = ast::NameRef::cast(expr.syntax().clone()) {
+        name_ref
+    } else {
+        ast::FieldExpr::cast(expr.syntax().clone())?.field()?
+    };
+
+    let function_ptr = resolve::resolve_name_ref(binder, &name_ref)?;
+
+    let root = file.syntax();
+    let function_name_node = function_ptr.to_node(root);
+
+    if let Some(create_function) = function_name_node
+        .ancestors()
+        .find_map(ast::CreateFunction::cast)
+        && let Some(param_list) = create_function.param_list()
+    {
+        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();
+                hints.push(InlayHint {
+                    position: arg_start,
+                    label: format!("{}: ", param_name.syntax().text()),
+                    kind: InlayHintKind::Parameter,
+                });
+            }
+        }
+    };
+
+    Some(())
+}
+
+#[cfg(test)]
+mod test {
+    use crate::inlay_hints::inlay_hints;
+    use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
+    use insta::assert_snapshot;
+    use squawk_syntax::ast;
+
+    #[track_caller]
+    fn check_inlay_hints(sql: &str) -> String {
+        let parse = ast::SourceFile::parse(sql);
+        assert_eq!(parse.errors(), vec![]);
+        let file: ast::SourceFile = parse.tree();
+
+        let hints = inlay_hints(&file);
+
+        if hints.is_empty() {
+            return String::new();
+        }
+
+        let mut modified_sql = sql.to_string();
+        let mut insertions: Vec<(usize, String)> = hints
+            .iter()
+            .map(|hint| {
+                let offset: usize = hint.position.into();
+                (offset, hint.label.clone())
+            })
+            .collect();
+
+        insertions.sort_by(|a, b| b.0.cmp(&a.0));
+
+        for (offset, label) in &insertions {
+            modified_sql.insert_str(*offset, label);
+        }
+
+        let mut annotations = vec![];
+        let mut cumulative_offset = 0;
+
+        insertions.reverse();
+        for (original_offset, label) in insertions {
+            let new_offset = original_offset + cumulative_offset;
+            annotations.push((new_offset, label.len()));
+            cumulative_offset += label.len();
+        }
+
+        let mut snippet = Snippet::source(&modified_sql).fold(true);
+
+        for (offset, len) in annotations {
+            snippet = snippet.annotation(AnnotationKind::Context.span(offset..offset + len));
+        }
+
+        let group = Level::INFO.primary_title("inlay hints").element(snippet);
+
+        let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
+        renderer
+            .render(&[group])
+            .to_string()
+            .replace("info: inlay hints", "inlay hints:")
+    }
+
+    #[test]
+    fn single_param() {
+        assert_snapshot!(check_inlay_hints("
+create function foo(a int) returns int as 'select $$1' language sql;
+select foo(1);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ select foo(a: 1);
+          ╰╴           ───
+        ");
+    }
+
+    #[test]
+    fn multiple_params() {
+        assert_snapshot!(check_inlay_hints("
+create function add(a int, b int) returns int as 'select $$1 + $$2' language sql;
+select add(1, 2);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ select add(a: 1, b: 2);
+          ╰╴           ───   ───
+        ");
+    }
+
+    #[test]
+    fn no_params() {
+        assert_snapshot!(check_inlay_hints("
+create function foo() returns int as 'select 1' language sql;
+select foo();
+"), @"");
+    }
+
+    #[test]
+    fn with_schema() {
+        assert_snapshot!(check_inlay_hints("
+create function public.foo(x int) returns int as 'select $$1' language sql;
+select public.foo(42);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ select public.foo(x: 42);
+          ╰╴                  ───
+        ");
+    }
+
+    #[test]
+    fn with_search_path() {
+        assert_snapshot!(check_inlay_hints(r#"
+set search_path to myschema;
+create function foo(val int) returns int as 'select $$1' language sql;
+select foo(100);
+"#), @r"
+        inlay hints:
+          ╭▸ 
+        4 │ select foo(val: 100);
+          ╰╴           ─────
+        ");
+    }
+
+    #[test]
+    fn multiple_calls() {
+        assert_snapshot!(check_inlay_hints("
+create function inc(n int) returns int as 'select $$1 + 1' language sql;
+select inc(1), inc(2);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ select inc(n: 1), inc(n: 2);
+          ╰╴           ───        ───
+        ");
+    }
+
+    #[test]
+    fn more_args_than_params() {
+        assert_snapshot!(check_inlay_hints("
+create function foo(a int) returns int as 'select $$1' language sql;
+select foo(1, 2);
+"), @r"
+        inlay hints:
+          ╭▸ 
+        3 │ select foo(a: 1, 2);
+          ╰╴           ───
+        ");
+    }
+}

crates/squawk_ide/src/lib.rs

@@ -6,6 +6,7 @@ pub mod find_references;
 mod generated;
 pub mod goto_definition;
 pub mod hover;
+pub mod inlay_hints;
 mod offsets;
 mod resolve;
 mod scope;