Skip to content

Commit c05b965

Browse files
committed
fix(lsp): preserve annotated type-variable names when adding annotations
``` fn wibble(a, b, c) { ↑ let x: a = todo fn(a: b, b: c) -> d { todo } } ``` After action: ```diff - fn wibble(a: e, b: f, c: g) -> fn(b, c) -> h + fn wibble(a: e, b: f, c: g) -> fn(b, c) -> d ^ original name is preserved ``` When the "add annotations" code-action is used, the compiler now retains the original type-variable names from the user's annotation instead of replacing them with fresh or default names. This ensures that added annotations reflect the user's naming exactly and prevents confusing or misleading variable renaming.
1 parent d46fbbc commit c05b965

3 files changed

Lines changed: 181 additions & 1 deletion

File tree

compiler-core/src/language_server/code_action.rs

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1372,13 +1372,187 @@ fn collect_type_variables(printer: &mut Printer<'_>, function: &ast::TypedFuncti
13721372
}
13731373

13741374
impl<'ast, 'a, 'b> ast::visit::Visit<'ast> for TypeVariableCollector<'a, 'b> {
1375+
fn visit_typed_function(&mut self, fun: &'ast ast::TypedFunction) {
1376+
for argument in fun.arguments.iter() {
1377+
if let Some(annotation) = &argument.annotation {
1378+
register_type_variables_from_annotation(
1379+
self.printer,
1380+
annotation,
1381+
argument.type_.as_ref(),
1382+
);
1383+
}
1384+
}
1385+
1386+
if let Some(annotation) = &fun.return_annotation {
1387+
register_type_variables_from_annotation(
1388+
self.printer,
1389+
annotation,
1390+
fun.return_type.as_ref(),
1391+
);
1392+
}
1393+
1394+
ast::visit::visit_typed_function(self, fun);
1395+
}
1396+
1397+
fn visit_typed_expr_fn(
1398+
&mut self,
1399+
location: &'ast SrcSpan,
1400+
type_: &'ast Arc<Type>,
1401+
kind: &'ast FunctionLiteralKind,
1402+
arguments: &'ast [TypedArg],
1403+
body: &'ast Vec1<TypedStatement>,
1404+
return_annotation: &'ast Option<ast::TypeAst>,
1405+
) {
1406+
if let Type::Fn {
1407+
arguments: argument_types,
1408+
return_: return_type,
1409+
..
1410+
} = type_.as_ref()
1411+
{
1412+
for (argument, argument_type) in arguments.iter().zip(argument_types) {
1413+
if let Some(annotation) = &argument.annotation {
1414+
register_type_variables_from_annotation(
1415+
self.printer,
1416+
annotation,
1417+
argument_type.as_ref(),
1418+
);
1419+
}
1420+
}
1421+
1422+
if let Some(annotation) = return_annotation {
1423+
register_type_variables_from_annotation(
1424+
self.printer,
1425+
annotation,
1426+
return_type.as_ref(),
1427+
);
1428+
}
1429+
}
1430+
1431+
ast::visit::visit_typed_expr_fn(
1432+
self,
1433+
location,
1434+
type_,
1435+
kind,
1436+
arguments,
1437+
body,
1438+
return_annotation,
1439+
);
1440+
}
1441+
13751442
fn visit_type_ast_var(&mut self, _location: &'ast SrcSpan, name: &'ast EcoString) {
13761443
// Register this type variable so that we don't duplicate names when
13771444
// adding annotations.
13781445
self.printer.register_type_variable(name.clone());
13791446
}
13801447
}
13811448

1449+
fn register_type_variables_from_annotation(
1450+
printer: &mut Printer<'_>,
1451+
annotation: &ast::TypeAst,
1452+
type_: &Type,
1453+
) {
1454+
// fn wibble(a, b, c) {
1455+
// fn(a: b, b: c) -> d { ... }
1456+
// ^
1457+
// Without this tracking the printer could rename `d` to a fresh `h`.
1458+
match (annotation, type_) {
1459+
(ast::TypeAst::Var(ast::TypeAstVar { name, .. }), Type::Var { type_ }) => {
1460+
match &*type_.borrow() {
1461+
TypeVar::Generic { id } | TypeVar::Unbound { id } => {
1462+
let id = *id;
1463+
printer.register_type_variable(name.clone());
1464+
printer.register_type_variable_with_id(id, name.clone());
1465+
}
1466+
TypeVar::Link { type_ } => {
1467+
register_type_variables_from_annotation(printer, annotation, type_.as_ref());
1468+
}
1469+
}
1470+
}
1471+
1472+
(
1473+
ast::TypeAst::Fn(ast::TypeAstFn {
1474+
arguments: annotation_arguments,
1475+
return_: annotation_return,
1476+
..
1477+
}),
1478+
Type::Fn {
1479+
arguments: type_arguments,
1480+
return_: type_return,
1481+
..
1482+
},
1483+
) => {
1484+
for (argument_annotation, argument_type) in
1485+
annotation_arguments.iter().zip(type_arguments)
1486+
{
1487+
// Maintain the names from each `fn(arg: name, ...)` position.
1488+
register_type_variables_from_annotation(
1489+
printer,
1490+
argument_annotation,
1491+
argument_type.as_ref(),
1492+
);
1493+
}
1494+
1495+
// And likewise propagate the annotated return variable.
1496+
register_type_variables_from_annotation(
1497+
printer,
1498+
annotation_return.as_ref(),
1499+
type_return.as_ref(),
1500+
);
1501+
}
1502+
1503+
(
1504+
ast::TypeAst::Constructor(ast::TypeAstConstructor {
1505+
arguments: annotation_arguments,
1506+
..
1507+
}),
1508+
Type::Named {
1509+
arguments: type_arguments,
1510+
..
1511+
},
1512+
) => {
1513+
for (argument_annotation, argument_type) in
1514+
annotation_arguments.iter().zip(type_arguments)
1515+
{
1516+
// Track aliases introduced inside named type arguments.
1517+
register_type_variables_from_annotation(
1518+
printer,
1519+
argument_annotation,
1520+
argument_type.as_ref(),
1521+
);
1522+
}
1523+
}
1524+
1525+
(
1526+
ast::TypeAst::Tuple(ast::TypeAstTuple {
1527+
elements: annotation_elements,
1528+
..
1529+
}),
1530+
Type::Tuple {
1531+
elements: type_elements,
1532+
..
1533+
},
1534+
) => {
1535+
for (element_annotation, element_type) in annotation_elements.iter().zip(type_elements)
1536+
{
1537+
// Tuples can hide extra annotations; ensure each slot retains its label.
1538+
register_type_variables_from_annotation(
1539+
printer,
1540+
element_annotation,
1541+
element_type.as_ref(),
1542+
);
1543+
}
1544+
}
1545+
1546+
(_, Type::Var { type_ }) => {
1547+
if let TypeVar::Link { type_ } = &*type_.borrow() {
1548+
register_type_variables_from_annotation(printer, annotation, type_.as_ref());
1549+
}
1550+
}
1551+
1552+
_ => {}
1553+
}
1554+
}
1555+
13821556
pub struct QualifiedConstructor<'a> {
13831557
import: &'a Import<EcoString>,
13841558
used_name: EcoString,

compiler-core/src/language_server/tests/snapshots/gleam_core__language_server__tests__action__type_variables_in_let_bindings_are_considered_when_adding_annotations.snap

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ fn wibble(a, b, c) {
1515

1616
----- AFTER ACTION
1717

18-
fn wibble(a: e, b: f, c: g) -> fn(b, c) -> h {
18+
fn wibble(a: e, b: f, c: g) -> fn(b, c) -> d {
1919
let x: a = todo
2020
fn(a: b, b: c) -> d {
2121
todo

compiler-core/src/type_/printer.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,12 @@ impl<'a> Printer<'a> {
454454
_ = self.printed_type_variable_names.insert(name);
455455
}
456456

457+
/// Record that the given type-variable id is already using the supplied name.
458+
pub fn register_type_variable_with_id(&mut self, id: u64, name: EcoString) {
459+
_ = self.printed_type_variable_names.insert(name.clone());
460+
_ = self.printed_type_variables.insert(id, name);
461+
}
462+
457463
pub fn print_type(&mut self, type_: &Type) -> EcoString {
458464
let mut buffer = EcoString::new();
459465
self.print(type_, &mut buffer, PrintMode::Normal);

0 commit comments

Comments
 (0)