Skip to main content

gruel_lsp/
inlay_hints.rs

1//! Inlay hints (ADR-0091 Phase 6).
2//!
3//! - After each `let` binding without an explicit type annotation, show
4//!   the inferred type (`: i32`).
5//! - After each unnamed call argument, show the parameter name (`x: 42`).
6
7use gruel_compiler::{Type, TypeInternPool};
8use gruel_parser::ast::{Ast, BlockExpr, Expr, Function, Item, Pattern, Statement};
9use gruel_util::{FileId, Span};
10use lasso::{Spur, ThreadedRodeo};
11use rustc_hash::FxHashMap;
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum InlayKind {
15    Type,
16    Parameter,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct InlayHint {
21    pub label: String,
22    pub kind: InlayKind,
23    /// Position (as byte offset) where the hint should render.
24    pub byte: u32,
25    /// FileId the byte is within.
26    pub file_id: FileId,
27}
28
29/// Produce all inlay hints for a file. `expr_types` is the side-table
30/// populated by Phase 4.
31pub fn inlay_hints(
32    ast: &Ast,
33    interner: &ThreadedRodeo,
34    expr_types: &FxHashMap<Span, Type>,
35    type_pool: Option<&TypeInternPool>,
36    file_id: FileId,
37) -> Vec<InlayHint> {
38    let mut hints = Vec::new();
39    for item in &ast.items {
40        match item {
41            Item::Function(f) => {
42                if f.span.file_id == file_id {
43                    visit_expr(
44                        &f.body,
45                        interner,
46                        expr_types,
47                        type_pool,
48                        file_id,
49                        Some(ast),
50                        &mut hints,
51                    );
52                }
53            }
54            Item::Struct(s) => {
55                for m in &s.methods {
56                    visit_expr(
57                        &m.body,
58                        interner,
59                        expr_types,
60                        type_pool,
61                        file_id,
62                        Some(ast),
63                        &mut hints,
64                    );
65                }
66            }
67            Item::Enum(e) => {
68                for m in &e.methods {
69                    visit_expr(
70                        &m.body,
71                        interner,
72                        expr_types,
73                        type_pool,
74                        file_id,
75                        Some(ast),
76                        &mut hints,
77                    );
78                }
79            }
80            Item::Derive(d) => {
81                for m in &d.methods {
82                    visit_expr(
83                        &m.body,
84                        interner,
85                        expr_types,
86                        type_pool,
87                        file_id,
88                        Some(ast),
89                        &mut hints,
90                    );
91                }
92            }
93            _ => {}
94        }
95    }
96    hints
97}
98
99fn visit_expr(
100    expr: &Expr,
101    interner: &ThreadedRodeo,
102    expr_types: &FxHashMap<Span, Type>,
103    type_pool: Option<&TypeInternPool>,
104    file_id: FileId,
105    ast: Option<&Ast>,
106    out: &mut Vec<InlayHint>,
107) {
108    match expr {
109        Expr::Block(b) => visit_block(b, interner, expr_types, type_pool, file_id, ast, out),
110        Expr::Call(c) => {
111            // Look up the callee function for argument hints.
112            if let Some(ast) = ast {
113                if let Some(callee) = find_function(ast, c.name.name) {
114                    for (i, arg) in c.args.iter().enumerate() {
115                        if let Some(p) = callee.params.get(i) {
116                            // Only suggest when the arg is a bare literal
117                            // (i.e. it's not already an identifier matching
118                            // the param name).
119                            if matches!(arg.expr, Expr::Ident(id) if id.name == p.name.name) {
120                                continue;
121                            }
122                            out.push(InlayHint {
123                                label: format!("{}:", interner.resolve(&p.name.name)),
124                                kind: InlayKind::Parameter,
125                                byte: arg.span.start,
126                                file_id: arg.span.file_id,
127                            });
128                        }
129                    }
130                }
131            }
132            for arg in &c.args {
133                visit_expr(
134                    &arg.expr, interner, expr_types, type_pool, file_id, ast, out,
135                );
136            }
137        }
138        Expr::If(i) => {
139            visit_expr(&i.cond, interner, expr_types, type_pool, file_id, ast, out);
140            visit_block(
141                &i.then_block,
142                interner,
143                expr_types,
144                type_pool,
145                file_id,
146                ast,
147                out,
148            );
149            if let Some(b) = &i.else_block {
150                visit_block(b, interner, expr_types, type_pool, file_id, ast, out);
151            }
152        }
153        Expr::While(w) => {
154            visit_expr(&w.cond, interner, expr_types, type_pool, file_id, ast, out);
155            visit_block(&w.body, interner, expr_types, type_pool, file_id, ast, out);
156        }
157        Expr::For(f) => {
158            visit_expr(
159                &f.iterable,
160                interner,
161                expr_types,
162                type_pool,
163                file_id,
164                ast,
165                out,
166            );
167            visit_block(&f.body, interner, expr_types, type_pool, file_id, ast, out);
168        }
169        Expr::Loop(l) => visit_block(&l.body, interner, expr_types, type_pool, file_id, ast, out),
170        Expr::Match(m) => {
171            visit_expr(
172                &m.scrutinee,
173                interner,
174                expr_types,
175                type_pool,
176                file_id,
177                ast,
178                out,
179            );
180            for arm in &m.arms {
181                visit_expr(
182                    &arm.body, interner, expr_types, type_pool, file_id, ast, out,
183                );
184            }
185        }
186        Expr::Binary(b) => {
187            visit_expr(&b.left, interner, expr_types, type_pool, file_id, ast, out);
188            visit_expr(&b.right, interner, expr_types, type_pool, file_id, ast, out);
189        }
190        Expr::Unary(u) => visit_expr(
191            &u.operand, interner, expr_types, type_pool, file_id, ast, out,
192        ),
193        Expr::Paren(p) => visit_expr(&p.inner, interner, expr_types, type_pool, file_id, ast, out),
194        Expr::Return(r) => {
195            if let Some(e) = &r.value {
196                visit_expr(e, interner, expr_types, type_pool, file_id, ast, out);
197            }
198        }
199        Expr::Tuple(t) => {
200            for e in &t.elems {
201                visit_expr(e, interner, expr_types, type_pool, file_id, ast, out);
202            }
203        }
204        Expr::Index(i) => {
205            visit_expr(&i.base, interner, expr_types, type_pool, file_id, ast, out);
206            visit_expr(&i.index, interner, expr_types, type_pool, file_id, ast, out);
207        }
208        Expr::MethodCall(m) => {
209            visit_expr(
210                &m.receiver,
211                interner,
212                expr_types,
213                type_pool,
214                file_id,
215                ast,
216                out,
217            );
218            for arg in &m.args {
219                visit_expr(
220                    &arg.expr, interner, expr_types, type_pool, file_id, ast, out,
221                );
222            }
223        }
224        _ => {}
225    }
226}
227
228fn visit_block(
229    b: &BlockExpr,
230    interner: &ThreadedRodeo,
231    expr_types: &FxHashMap<Span, Type>,
232    type_pool: Option<&TypeInternPool>,
233    file_id: FileId,
234    ast: Option<&Ast>,
235    out: &mut Vec<InlayHint>,
236) {
237    for stmt in &b.statements {
238        if let Statement::Let(l) = stmt {
239            if l.ty.is_none() {
240                if let Pattern::Ident { name, .. } = &l.pattern {
241                    if name.span.file_id == file_id {
242                        if let Some(ty) = lookup_init_type(&l.init, expr_types) {
243                            let label = if let Some(pool) = type_pool {
244                                format!(": {}", pool.format_type_name(ty))
245                            } else {
246                                format!(": {:?}", ty)
247                            };
248                            out.push(InlayHint {
249                                label,
250                                kind: InlayKind::Type,
251                                byte: name.span.end,
252                                file_id: name.span.file_id,
253                            });
254                        }
255                    }
256                }
257            }
258            visit_expr(&l.init, interner, expr_types, type_pool, file_id, ast, out);
259        } else if let Statement::Expr(e) = stmt {
260            visit_expr(e, interner, expr_types, type_pool, file_id, ast, out);
261        } else if let Statement::Assign(a) = stmt {
262            visit_expr(&a.value, interner, expr_types, type_pool, file_id, ast, out);
263        }
264    }
265    visit_expr(&b.expr, interner, expr_types, type_pool, file_id, ast, out);
266}
267
268fn lookup_init_type(init: &Expr, expr_types: &FxHashMap<Span, Type>) -> Option<Type> {
269    let span = expr_span(init)?;
270    // Try the exact init span first.
271    if let Some(ty) = expr_types.get(&span) {
272        return Some(*ty);
273    }
274    // Otherwise: any expr_types entry that exactly equals the init span.
275    None
276}
277
278fn expr_span(e: &Expr) -> Option<Span> {
279    match e {
280        Expr::Int(l) => Some(l.span),
281        Expr::Float(l) => Some(l.span),
282        Expr::String(l) => Some(l.span),
283        Expr::Char(l) => Some(l.span),
284        Expr::Bool(l) => Some(l.span),
285        Expr::Unit(l) => Some(l.span),
286        Expr::Ident(i) => Some(i.span),
287        Expr::Binary(b) => Some(b.span),
288        Expr::Unary(u) => Some(u.span),
289        Expr::Paren(p) => Some(p.span),
290        Expr::Block(b) => Some(b.span),
291        Expr::If(i) => Some(i.span),
292        Expr::Match(m) => Some(m.span),
293        Expr::While(w) => Some(w.span),
294        Expr::For(f) => Some(f.span),
295        Expr::Loop(l) => Some(l.span),
296        Expr::Call(c) => Some(c.span),
297        Expr::MethodCall(m) => Some(m.span),
298        Expr::Field(f) => Some(f.span),
299        Expr::Tuple(t) => Some(t.span),
300        Expr::Index(i) => Some(i.span),
301        Expr::Return(r) => Some(r.span),
302        _ => None,
303    }
304}
305
306fn find_function(ast: &Ast, name: Spur) -> Option<Function> {
307    for item in &ast.items {
308        if let Item::Function(f) = item {
309            if f.name.name == name {
310                return Some(f.clone());
311            }
312        }
313    }
314    None
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320    use gruel_compiler::PreviewFeatures;
321    use gruel_target::Target;
322
323    fn snap_for(source: &str) -> std::sync::Arc<crate::analysis::Snapshot> {
324        use crate::analysis::{WorkspaceFile, analyze};
325        use std::path::PathBuf;
326        let files = vec![WorkspaceFile {
327            path: PathBuf::from("main.gruel"),
328            text: source.to_string(),
329            file_id: FileId::new(1),
330        }];
331        let res = analyze(&files, &PreviewFeatures::default(), &Target::host());
332        std::sync::Arc::new(res.snapshot.unwrap())
333    }
334
335    #[test]
336    fn inlay_for_untyped_let() {
337        let src = "fn main() -> i32 { let answer = 42; answer }";
338        let snap = snap_for(src);
339        let hints = inlay_hints(
340            &snap.ast,
341            &snap.interner,
342            &snap.expr_types,
343            snap.type_pool.as_deref(),
344            FileId::new(1),
345        );
346        let type_hints: Vec<_> = hints.iter().filter(|h| h.kind == InlayKind::Type).collect();
347        assert!(!type_hints.is_empty(), "expected at least one type hint");
348        assert!(type_hints.iter().any(|h| h.label.contains("i32")));
349    }
350
351    #[test]
352    fn parameter_hint_for_call() {
353        let src = "fn add(x: i32, y: i32) -> i32 { x + y }\nfn main() -> i32 { add(1, 2) }";
354        let snap = snap_for(src);
355        let hints = inlay_hints(
356            &snap.ast,
357            &snap.interner,
358            &snap.expr_types,
359            snap.type_pool.as_deref(),
360            FileId::new(1),
361        );
362        let param_hints: Vec<_> = hints
363            .iter()
364            .filter(|h| h.kind == InlayKind::Parameter)
365            .collect();
366        assert!(
367            param_hints.iter().any(|h| h.label == "x:"),
368            "expected `x:` hint, got: {:?}",
369            hints
370        );
371    }
372}