1use crate::collect;
2use crate::db::{File, parse};
3use crate::file::InFile;
4use crate::goto_definition;
5use crate::resolve;
6use crate::symbols::Name;
7use rowan::{TextRange, TextSize};
8use salsa::Database as Db;
9use squawk_syntax::ast::{self, AstNode};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum InlayHintKind {
14 Type,
15 Parameter,
16}
17
18#[derive(Clone, PartialEq, Eq)]
19pub struct InlayHint {
20 pub position: TextSize,
21 pub label: String,
22 pub kind: InlayHintKind,
23 pub target: Option<InFile<TextRange>>,
26}
27
28#[salsa::tracked]
29pub fn inlay_hints(db: &dyn Db, file: File) -> Vec<InlayHint> {
30 let mut hints = vec![];
31 for node in parse(db, file).tree().syntax().descendants() {
32 if let Some(call_expr) = ast::CallExpr::cast(node.clone()) {
33 inlay_hint_call_expr(db, &mut hints, file, call_expr);
34 } else if let Some(insert) = ast::Insert::cast(node) {
35 inlay_hint_insert(db, &mut hints, file, insert);
36 }
37 }
38 hints
39}
40
41fn inlay_hint_call_expr(
42 db: &dyn Db,
43 hints: &mut Vec<InlayHint>,
44 file_id: File,
45 call_expr: ast::CallExpr,
46) -> Option<()> {
47 let arg_list = call_expr.arg_list()?;
48 let expr = call_expr.expr()?;
49
50 let name_ref = if let Some(name_ref) = ast::NameRef::cast(expr.syntax().clone()) {
51 name_ref
52 } else {
53 ast::FieldExpr::cast(expr.syntax().clone())?.field()?
54 };
55
56 let location = goto_definition::goto_definition(
57 db,
58 InFile::new(file_id, name_ref.syntax().text_range().start()),
59 )
60 .into_iter()
61 .next()?;
62
63 let def_file = parse(db, location.file).tree();
64
65 let function_name_node = def_file.syntax().covering_element(location.range);
66
67 if let Some(create_function) = function_name_node
68 .ancestors()
69 .find_map(ast::CreateFunction::cast)
70 && let Some(param_list) = create_function.param_list()
71 {
72 for (param, arg) in param_list.params().zip(arg_list.args()) {
73 if let Some(param_name) = param.name() {
74 let arg_start = arg.syntax().text_range().start();
75 let target = Some(InFile::new(location.file, param_name.syntax().text_range()));
76 hints.push(InlayHint {
77 position: arg_start,
78 label: format!("{}: ", param_name.syntax().text()),
79 kind: InlayHintKind::Parameter,
80 target,
81 });
82 }
83 }
84 };
85
86 Some(())
87}
88
89fn inlay_hint_insert(
90 db: &dyn Db,
91 hints: &mut Vec<InlayHint>,
92 file_id: File,
93 insert: ast::Insert,
94) -> Option<()> {
95 let name_start = insert
96 .path()?
97 .segment()?
98 .name_ref()?
99 .syntax()
100 .text_range()
101 .start();
102 let location = goto_definition::goto_definition(db, InFile::new(file_id, name_start))
105 .into_iter()
106 .next();
107
108 let def_file = location.as_ref().map(|loc| loc.file).unwrap_or(file_id);
109 let def_tree = parse(db, def_file).tree();
110
111 let create_table = location.as_ref().and_then(|loc| {
112 def_tree
113 .syntax()
114 .covering_element(loc.range)
115 .ancestors()
116 .find_map(ast::CreateTableLike::cast)
117 });
118
119 let columns: Vec<(Name, Option<InFile<TextRange>>)> =
120 if let Some(column_list) = insert.column_list() {
121 column_list
123 .columns()
124 .filter_map(|col| {
125 let col_name = col.name_ref().map(|x| Name::from_node(&x))?;
126 let target = create_table
127 .as_ref()
128 .and_then(|x| {
129 resolve::find_column_in_create_table(
130 db,
131 InFile::new(def_file, x),
132 &col_name,
133 )
134 })
135 .and_then(|x| x.into_iter().next())
136 .map(|x| InFile::new(x.file, x.range));
137 Some((col_name, target))
138 })
139 .collect()
140 } else {
141 collect::columns_from_create_table(db, def_file, &create_table?)
143 .into_iter()
144 .map(|(col_name, ptr)| {
145 let target = ptr.map(|ptr| InFile::new(ptr.file_id, ptr.value.text_range()));
146 (col_name, target)
147 })
148 .collect()
149 };
150
151 inlay_hint_insert_select(hints, columns, insert.select_variant()?)
152}
153
154fn inlay_hint_insert_select(
155 hints: &mut Vec<InlayHint>,
156 columns: Vec<(Name, Option<InFile<TextRange>>)>,
157 select_variant: ast::SelectVariant,
158) -> Option<()> {
159 if let ast::SelectVariant::Values(values) = &select_variant {
160 for row in values.row_list()?.rows() {
162 for ((column_name, target), expr) in columns.iter().zip(row.exprs()) {
163 let expr_start = expr.syntax().text_range().start();
164 hints.push(InlayHint {
165 position: expr_start,
166 label: format!("{column_name}: "),
167 kind: InlayHintKind::Parameter,
168 target: *target,
169 });
170 }
171 }
172 return Some(());
173 }
174
175 let target_list = select_variant.target_list()?;
177 for ((column_name, target), target_expr) in columns.iter().zip(target_list.targets()) {
178 let expr = target_expr.expr()?;
179 let expr_start = expr.syntax().text_range().start();
180 hints.push(InlayHint {
181 position: expr_start,
182 label: format!("{column_name}: "),
183 kind: InlayHintKind::Parameter,
184 target: *target,
185 });
186 }
187
188 Some(())
189}
190
191#[cfg(test)]
192mod test {
193 use crate::builtins::builtins_file;
194 use crate::db::{Database, File};
195 use crate::inlay_hints::{InlayHint, inlay_hints};
196 use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
197 use insta::assert_snapshot;
198 use rustc_hash::FxHashMap;
199 use std::ops::Range;
200
201 #[must_use]
202 #[track_caller]
203 fn check_inlay_hints(sql: &str) -> String {
204 let db = Database::default();
205 let file = File::new(&db, sql.to_string().into());
206
207 assert_eq!(crate::db::parse(&db, file).errors(), vec![]);
208
209 let hints = inlay_hints(&db, file);
210
211 if hints.is_empty() {
212 return String::new();
213 }
214
215 let mut modified_sql = sql.to_string();
216 let mut indexed: Vec<(usize, &InlayHint)> = hints.iter().enumerate().collect();
217 indexed.sort_by_key(|(_, h)| h.position);
218
219 let mut label_annotations: Vec<Range<usize>> = vec![0..0; hints.len()];
220 let mut cumulative = 0;
221 for (i, hint) in &indexed {
222 let pos: usize = hint.position.into();
223 let new_pos = pos + cumulative;
224 modified_sql.insert_str(new_pos, &hint.label);
225 label_annotations[*i] = new_pos..new_pos + hint.label.len();
226 cumulative += hint.label.len();
227 }
228
229 let mut targets_by_file: FxHashMap<File, Vec<(usize, Range<usize>)>> = FxHashMap::default();
230 for (i, hint) in hints.iter().enumerate() {
231 if let Some(target) = &hint.target {
232 let start: usize = target.value.start().into();
233 let end: usize = target.value.end().into();
234 targets_by_file
235 .entry(target.file_id)
236 .or_default()
237 .push((i + 1, start..end));
238 }
239 }
240
241 let mut file_paths: FxHashMap<File, &'static str> = FxHashMap::default();
242 file_paths.insert(file, "current.sql");
243 file_paths.insert(builtins_file(&db), "builtins.sql");
244
245 let mut labels_snippet = Snippet::source(&modified_sql).fold(true);
246 for (i, range) in label_annotations.into_iter().enumerate() {
247 labels_snippet = labels_snippet.annotation(
248 AnnotationKind::Context
249 .span(range)
250 .label(format!("{}. label", i + 1)),
251 );
252 }
253
254 let mut groups = vec![Level::INFO.primary_title("labels").element(labels_snippet)];
255
256 let mut target_entries = targets_by_file.into_iter().collect::<Vec<_>>();
257 target_entries.sort_by_key(|(_, targets)| {
258 targets.iter().map(|(i, _)| *i).min().unwrap_or(usize::MAX)
259 });
260
261 let target_contents = target_entries
262 .into_iter()
263 .map(|(f, targets)| {
264 let path = *file_paths.get(&f).unwrap();
265 (f.content(&db).clone(), path, targets)
266 })
267 .collect::<Vec<_>>();
268
269 for (content, path, targets) in &target_contents {
270 let mut snippet = Snippet::source(content.as_ref()).fold(true).path(*path);
271 for (i, range) in targets {
272 snippet = snippet.annotation(
273 AnnotationKind::Context
274 .span(range.clone())
275 .label(format!("{i}. target")),
276 );
277 }
278 groups.push(Level::INFO.primary_title("targets").element(snippet));
279 }
280
281 let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
282 renderer
283 .render(&groups)
284 .to_string()
285 .replace("info: labels", "labels:")
286 .replace("info: targets", "targets:")
287 }
288
289 #[test]
290 fn single_param() {
291 assert_snapshot!(check_inlay_hints("
292create function foo(a int) returns int as 'select $$1' language sql;
293select foo(1);
294"), @"
295 labels:
296 ╭▸
297 3 │ select foo(a: 1);
298 │ ─── 1. label
299 ╰╴
300 targets:
301 ╭▸ current.sql:2:21
302 │
303 2 │ create function foo(a int) returns int as 'select $$1' language sql;
304 ╰╴ ─ 1. target
305 ");
306 }
307
308 #[test]
309 fn multiple_params() {
310 assert_snapshot!(check_inlay_hints("
311create function add(a int, b int) returns int as 'select $$1 + $$2' language sql;
312select add(1, 2);
313"), @"
314 labels:
315 ╭▸
316 3 │ select add(a: 1, b: 2);
317 │ ┬── ─── 2. label
318 │ │
319 │ 1. label
320 ╰╴
321 targets:
322 ╭▸ current.sql:2:21
323 │
324 2 │ create function add(a int, b int) returns int as 'select $$1 + $$2' language sql;
325 │ ┬ ─ 2. target
326 │ │
327 ╰╴ 1. target
328 ");
329 }
330
331 #[test]
332 fn no_params() {
333 assert_snapshot!(check_inlay_hints("
334create function foo() returns int as 'select 1' language sql;
335select foo();
336"), @"");
337 }
338
339 #[test]
340 fn with_schema() {
341 assert_snapshot!(check_inlay_hints("
342create function public.foo(x int) returns int as 'select $$1' language sql;
343select public.foo(42);
344"), @"
345 labels:
346 ╭▸
347 3 │ select public.foo(x: 42);
348 │ ─── 1. label
349 ╰╴
350 targets:
351 ╭▸ current.sql:2:28
352 │
353 2 │ create function public.foo(x int) returns int as 'select $$1' language sql;
354 ╰╴ ─ 1. target
355 ");
356 }
357
358 #[test]
359 fn with_search_path() {
360 assert_snapshot!(check_inlay_hints(r#"
361set search_path to myschema;
362create function foo(val int) returns int as 'select $$1' language sql;
363select foo(100);
364"#), @"
365 labels:
366 ╭▸
367 4 │ select foo(val: 100);
368 │ ───── 1. label
369 ╰╴
370 targets:
371 ╭▸ current.sql:3:21
372 │
373 3 │ create function foo(val int) returns int as 'select $$1' language sql;
374 ╰╴ ─── 1. target
375 ");
376 }
377
378 #[test]
379 fn multiple_calls() {
380 assert_snapshot!(check_inlay_hints("
381create function inc(n int) returns int as 'select $$1 + 1' language sql;
382select inc(1), inc(2);
383"), @"
384 labels:
385 ╭▸
386 3 │ select inc(n: 1), inc(n: 2);
387 │ ┬── ─── 2. label
388 │ │
389 │ 1. label
390 ╰╴
391 targets:
392 ╭▸ current.sql:2:21
393 │
394 2 │ create function inc(n int) returns int as 'select $$1 + 1' language sql;
395 │ ┬
396 │ │
397 │ 1. target
398 ╰╴ 2. target
399 ");
400 }
401
402 #[test]
403 fn more_args_than_params() {
404 assert_snapshot!(check_inlay_hints("
405create function foo(a int) returns int as 'select $$1' language sql;
406select foo(1, 2);
407"), @"
408 labels:
409 ╭▸
410 3 │ select foo(a: 1, 2);
411 │ ─── 1. label
412 ╰╴
413 targets:
414 ╭▸ current.sql:2:21
415 │
416 2 │ create function foo(a int) returns int as 'select $$1' language sql;
417 ╰╴ ─ 1. target
418 ");
419 }
420
421 #[test]
422 fn builtin_function() {
423 assert_snapshot!(check_inlay_hints("
424select json_strip_nulls('[1, null]', true);
425"), @"
426 labels:
427 ╭▸
428 2 │ select json_strip_nulls(target: '[1, null]', strip_in_arrays: true);
429 │ ──────── 1. label ───────────────── 2. label
430 ╰╴
431 targets:
432 ╭▸ builtins.sql:9239:45
433 │
434 9239 │ create function pg_catalog.json_strip_nulls(target json, strip_in_arrays boolean DEFAULT false) returns json
435 │ ┬───── ─────────────── 2. target
436 │ │
437 ╰╴ 1. target
438 ");
439 }
440
441 #[test]
442 fn insert_with_column_list() {
443 assert_snapshot!(check_inlay_hints("
444create table t (column_a int, column_b int, column_c text);
445insert into t (column_a, column_c) values (1, 'foo');
446"), @"
447 labels:
448 ╭▸
449 3 │ insert into t (column_a, column_c) values (column_a: 1, column_c: 'foo');
450 │ ┬───────── ────────── 2. label
451 │ │
452 │ 1. label
453 ╰╴
454 targets:
455 ╭▸ current.sql:2:17
456 │
457 2 │ create table t (column_a int, column_b int, column_c text);
458 ╰╴ ──────── 1. target ──────── 2. target
459 ");
460 }
461
462 #[test]
463 fn insert_without_column_list() {
464 assert_snapshot!(check_inlay_hints("
465create table t (column_a int, column_b int, column_c text);
466insert into t values (1, 2, 'foo');
467"), @"
468 labels:
469 ╭▸
470 3 │ insert into t values (column_a: 1, column_b: 2, column_c: 'foo');
471 │ ┬───────── ┬───────── ────────── 3. label
472 │ │ │
473 │ │ 2. label
474 │ 1. label
475 ╰╴
476 targets:
477 ╭▸ current.sql:2:17
478 │
479 2 │ create table t (column_a int, column_b int, column_c text);
480 │ ┬─────── ┬─────── ──────── 3. target
481 │ │ │
482 │ │ 2. target
483 ╰╴ 1. target
484 ");
485 }
486
487 #[test]
488 fn insert_multiple_rows() {
489 assert_snapshot!(check_inlay_hints("
490create table t (x int, y int);
491insert into t values (1, 2), (3, 4);
492"), @"
493 labels:
494 ╭▸
495 3 │ insert into t values (x: 1, y: 2), (x: 3, y: 4);
496 │ ┬── ┬── ┬── ─── 4. label
497 │ │ │ │
498 │ │ │ 3. label
499 │ │ 2. label
500 │ 1. label
501 ╰╴
502 targets:
503 ╭▸ current.sql:2:17
504 │
505 2 │ create table t (x int, y int);
506 │ ┬ ┬
507 │ │ │
508 │ │ 2. target
509 │ │ 4. target
510 │ 1. target
511 ╰╴ 3. target
512 ");
513 }
514
515 #[test]
516 fn insert_no_create_table() {
517 assert_snapshot!(check_inlay_hints("
518insert into t (a, b) values (1, 2);
519"), @"
520 labels:
521 ╭▸
522 2 │ insert into t (a, b) values (a: 1, b: 2);
523 │ ┬── ─── 2. label
524 │ │
525 ╰╴ 1. label
526 ");
527 }
528
529 #[test]
530 fn insert_more_values_than_columns() {
531 assert_snapshot!(check_inlay_hints("
532create table t (a int, b int);
533insert into t values (1, 2, 3);
534"), @"
535 labels:
536 ╭▸
537 3 │ insert into t values (a: 1, b: 2, 3);
538 │ ┬── ─── 2. label
539 │ │
540 │ 1. label
541 ╰╴
542 targets:
543 ╭▸ current.sql:2:17
544 │
545 2 │ create table t (a int, b int);
546 │ ┬ ─ 2. target
547 │ │
548 ╰╴ 1. target
549 ");
550 }
551
552 #[test]
553 fn insert_table_inherits_select() {
554 assert_snapshot!(check_inlay_hints("
555create table t (a int, b int);
556create table u (c int) inherits (t);
557insert into u select 1, 2, 3;
558"), @"
559 labels:
560 ╭▸
561 4 │ insert into u select a: 1, b: 2, c: 3;
562 │ ┬── ┬── ─── 3. label
563 │ │ │
564 │ │ 2. label
565 │ 1. label
566 ╰╴
567 targets:
568 ╭▸ current.sql:2:17
569 │
570 2 │ create table t (a int, b int);
571 │ ┬ ─ 2. target
572 │ │
573 │ 1. target
574 3 │ create table u (c int) inherits (t);
575 ╰╴ ─ 3. target
576 ");
577 }
578
579 #[test]
580 fn insert_table_inherits_builtin_values() {
581 assert_snapshot!(check_inlay_hints("
582create table t ()
583inherits (information_schema.sql_features);
584insert into t values (1, 2, 3, 4, 5, 6, 7);
585"), @"
586 labels:
587 ╭▸
588 4 │ …ues (feature_id: 1, feature_name: 2, sub_feature_id: 3, sub_feature_name: 4, is_supported: 5, is_verified_by: 6, comments: 7);
589 │ ┬─────────── ┬───────────── ┬─────────────── ┬───────────────── ┬───────────── ┬─────────────── ────────── 7. label
590 │ │ │ │ │ │ │
591 │ │ │ │ │ │ 6. label
592 │ │ │ │ │ 5. label
593 │ │ │ │ 4. label
594 │ │ │ 3. label
595 │ │ 2. label
596 │ 1. label
597 ╰╴
598 targets:
599 ╭▸ builtins.sql:436:3
600 │
601 436 │ feature_id information_schema.character_data,
602 │ ────────── 1. target
603 437 │ feature_name information_schema.character_data,
604 │ ──────────── 2. target
605 438 │ sub_feature_id information_schema.character_data,
606 │ ────────────── 3. target
607 439 │ sub_feature_name information_schema.character_data,
608 │ ──────────────── 4. target
609 440 │ is_supported information_schema.yes_or_no,
610 │ ──────────── 5. target
611 441 │ is_verified_by information_schema.character_data,
612 │ ────────────── 6. target
613 442 │ comments information_schema.character_data
614 ╰╴ ──────── 7. target
615 ");
616 }
617
618 #[test]
619 fn insert_table_inherits_create_table_as_values() {
620 assert_snapshot!(check_inlay_hints("
621create table parent as select 1 a, 'x'::text b;
622create table child (c int) inherits (parent);
623insert into child values (1, 2, 3);
624"), @"
625 labels:
626 ╭▸
627 4 │ insert into child values (a: 1, b: 2, c: 3);
628 │ ┬── ┬── ─── 3. label
629 │ │ │
630 │ │ 2. label
631 │ 1. label
632 ╰╴
633 targets:
634 ╭▸ current.sql:3:21
635 │
636 3 │ create table child (c int) inherits (parent);
637 ╰╴ ─ 3. target
638 ");
639 }
640
641 #[test]
642 fn insert_table_inherits_create_table_as_select_star() {
643 assert_snapshot!(check_inlay_hints("
644create table base (a int, b text);
645create table parent as select * from base;
646create table child (c int) inherits (parent);
647insert into child values (1, 2, 3);
648"), @"
649 labels:
650 ╭▸
651 5 │ insert into child values (a: 1, b: 2, c: 3);
652 │ ┬── ┬── ─── 3. label
653 │ │ │
654 │ │ 2. label
655 │ 1. label
656 ╰╴
657 targets:
658 ╭▸ current.sql:4:21
659 │
660 4 │ create table child (c int) inherits (parent);
661 ╰╴ ─ 3. target
662 ");
663 }
664
665 #[test]
666 fn insert_table_like_select() {
667 assert_snapshot!(check_inlay_hints("
668create table x (a int, b int);
669create table y (c int, like x);
670insert into y select 1, 2, 3;
671"), @"
672 labels:
673 ╭▸
674 4 │ insert into y select c: 1, a: 2, b: 3;
675 │ ┬── ┬── ─── 3. label
676 │ │ │
677 │ │ 2. label
678 │ 1. label
679 ╰╴
680 targets:
681 ╭▸ current.sql:2:17
682 │
683 2 │ create table x (a int, b int);
684 │ ┬ ─ 3. target
685 │ │
686 │ 2. target
687 3 │ create table y (c int, like x);
688 ╰╴ ─ 1. target
689 ");
690 }
691
692 #[test]
693 fn insert_select() {
694 assert_snapshot!(check_inlay_hints("
695create table t (a int, b int);
696insert into t select 1, 2;
697"), @"
698 labels:
699 ╭▸
700 3 │ insert into t select a: 1, b: 2;
701 │ ┬── ─── 2. label
702 │ │
703 │ 1. label
704 ╰╴
705 targets:
706 ╭▸ current.sql:2:17
707 │
708 2 │ create table t (a int, b int);
709 │ ┬ ─ 2. target
710 │ │
711 ╰╴ 1. target
712 ");
713 }
714
715 #[test]
716 fn insert_table_like_builtin_values() {
717 assert_snapshot!(check_inlay_hints("
718create table t (like information_schema.sql_features);
719insert into t values (1, 2, 3, 4, 5, 6, 7);
720"), @"
721 labels:
722 ╭▸
723 3 │ …ues (feature_id: 1, feature_name: 2, sub_feature_id: 3, sub_feature_name: 4, is_supported: 5, is_verified_by: 6, comments: 7);
724 │ ┬─────────── ┬───────────── ┬─────────────── ┬───────────────── ┬───────────── ┬─────────────── ────────── 7. label
725 │ │ │ │ │ │ │
726 │ │ │ │ │ │ 6. label
727 │ │ │ │ │ 5. label
728 │ │ │ │ 4. label
729 │ │ │ 3. label
730 │ │ 2. label
731 │ 1. label
732 ╰╴
733 targets:
734 ╭▸ builtins.sql:436:3
735 │
736 436 │ feature_id information_schema.character_data,
737 │ ────────── 1. target
738 437 │ feature_name information_schema.character_data,
739 │ ──────────── 2. target
740 438 │ sub_feature_id information_schema.character_data,
741 │ ────────────── 3. target
742 439 │ sub_feature_name information_schema.character_data,
743 │ ──────────────── 4. target
744 440 │ is_supported information_schema.yes_or_no,
745 │ ──────────── 5. target
746 441 │ is_verified_by information_schema.character_data,
747 │ ────────────── 6. target
748 442 │ comments information_schema.character_data
749 ╰╴ ──────── 7. target
750 ");
751 }
752
753 #[test]
754 fn insert_table_like_select_into_values() {
755 assert_snapshot!(check_inlay_hints("
756select 1 a, 'x'::text b into parent;
757create table child (like parent);
758insert into child values (1, 2);
759"), @"
760 labels:
761 ╭▸
762 4 │ insert into child values (a: 1, b: 2);
763 │ ┬── ─── 2. label
764 │ │
765 ╰╴ 1. label
766 ");
767 }
768}