1use 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 pub byte: u32,
25 pub file_id: FileId,
27}
28
29pub 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 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 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 if let Some(ty) = expr_types.get(&span) {
272 return Some(*ty);
273 }
274 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}