Localise type variables names when adding annotations

Author: GearsDatapacks

compiler-core/src/ast/visit.rs

@@ -637,6 +637,14 @@ pub fn visit_typed_function<'a, V>(v: &mut V, fun: &'a TypedFunction)
 where
     V: Visit<'a> + ?Sized,
 {
+    for argument in fun.arguments.iter() {
+        if let Some(annotation) = &argument.annotation {
+            v.visit_type_ast(annotation);
+        }
+    }
+    if let Some(annotation) = &fun.return_annotation {
+        v.visit_type_ast(annotation);
+    }
     for stmt in &fun.body {
         v.visit_typed_statement(stmt);
     }

compiler-core/src/language_server/code_action.rs

@@ -1170,6 +1170,11 @@ impl<'ast> ast::visit::Visit<'ast> for AddAnnotations<'_> {
     }
 
     fn visit_typed_module_constant(&mut self, constant: &'ast TypedModuleConstant) {
+        // Since type variable names are local to definitions, any type variables
+        // in other parts of the module shouldn't affect what we print for the
+        // annotations of this constant.
+        self.printer.clear_type_variables();
+
         let code_action_range = self.edits.src_span_to_lsp_range(constant.location);
 
         // Only offer the code action if the cursor is over the statement
@@ -1189,6 +1194,14 @@ impl<'ast> ast::visit::Visit<'ast> for AddAnnotations<'_> {
     }
 
     fn visit_typed_function(&mut self, fun: &'ast ast::TypedFunction) {
+        self.printer.clear_type_variables();
+
+        // Since type variable names are local to definitions, any type variables
+        // in other parts of the module shouldn't affect what we print for the
+        // annotations of this functions. The only variables which cannot clash
+        // are ones defined in the signature of this function, which we register
+        // when we visit the parameters of this function below.
+
         ast::visit::visit_typed_function(self, fun);
 
         let code_action_range = self.edits.src_span_to_lsp_range(
@@ -1227,6 +1240,12 @@ impl<'ast> ast::visit::Visit<'ast> for AddAnnotations<'_> {
         }
     }
 
+    fn visit_type_ast_var(&mut self, _location: &'ast SrcSpan, name: &'ast EcoString) {
+        // Register this type variable so that we don't duplicate names when
+        // adding annotations.
+        self.printer.register_type_variable(name.clone());
+    }
+
     fn visit_typed_expr_fn(
         &mut self,
         location: &'ast SrcSpan,
@@ -1420,19 +1439,6 @@ impl<'ast> ast::visit::Visit<'ast> for QualifiedToUnqualifiedImportFirstPass<'as
         );
     }
 
-    fn visit_typed_function(&mut self, fun: &'ast ast::TypedFunction) {
-        for arg in &fun.arguments {
-            if let Some(annotation) = &arg.annotation {
-                self.visit_type_ast(annotation);
-            }
-        }
-
-        if let Some(return_annotation) = &fun.return_annotation {
-            self.visit_type_ast(return_annotation);
-        }
-        ast::visit::visit_typed_function(self, fun);
-    }
-
     fn visit_type_ast_constructor(
         &mut self,
         location: &'ast SrcSpan,
@@ -1659,19 +1665,6 @@ impl<'ast> ast::visit::Visit<'ast> for QualifiedToUnqualifiedImportSecondPass<'a
         );
     }
 
-    fn visit_typed_function(&mut self, fun: &'ast ast::TypedFunction) {
-        for arg in &fun.arguments {
-            if let Some(annotation) = &arg.annotation {
-                self.visit_type_ast(annotation);
-            }
-        }
-
-        if let Some(return_annotation) = &fun.return_annotation {
-            self.visit_type_ast(return_annotation);
-        }
-        ast::visit::visit_typed_function(self, fun);
-    }
-
     fn visit_type_ast_constructor(
         &mut self,
         location: &'ast SrcSpan,
@@ -1904,18 +1897,6 @@ impl<'ast> ast::visit::Visit<'ast> for UnqualifiedToQualifiedImportFirstPass<'as
         );
     }
 
-    fn visit_typed_function(&mut self, fun: &'ast ast::TypedFunction) {
-        for arg in &fun.arguments {
-            if let Some(annotation) = &arg.annotation {
-                self.visit_type_ast(annotation);
-            }
-        }
-
-        if let Some(return_annotation) = &fun.return_annotation {
-            self.visit_type_ast(return_annotation);
-        }
-        ast::visit::visit_typed_function(self, fun);
-    }
     fn visit_type_ast_constructor(
         &mut self,
         location: &'ast SrcSpan,

compiler-core/src/language_server/tests/action.rs

@@ -9307,7 +9307,6 @@ fn collapse_nested_case_does_not_remove_labels() {
   }
 }
 
-pub type Wibble {
   Wibble(field: Int, field2: String)
   Wobble
 }
@@ -9328,7 +9327,6 @@ fn collapse_nested_case_does_not_remove_labels_with_shorthand_syntax() {
         2 -> 4
         _ -> -1
       }
-
     Wobble -> todo
   }
 }
@@ -9466,3 +9464,47 @@ fn collapse_nested_case_combines_inner_and_outer_guards_and_adds_parentheses_whe
         find_position_of("first").to_selection()
     );
 }
+
+// https://github.com/gleam-lang/gleam/issues/3786
+#[test]
+fn type_variables_from_other_functions_do_not_change_annotations() {
+    assert_code_action!(
+        ADD_ANNOTATIONS,
+        "
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+fn pair(a, b) {
+  #(a, b)
+}
+",
+        find_position_of("pair").to_selection()
+    );
+}
+
+#[test]
+fn type_variables_from_other_functions_do_not_change_annotations_constant() {
+    assert_code_action!(
+        ADD_ANNOTATION,
+        "
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+const empty = []
+",
+        find_position_of("empty").to_selection()
+    );
+}
+
+#[test]
+fn type_variables_are_not_duplicated_when_adding_annotations() {
+    assert_code_action!(
+        ADD_ANNOTATIONS,
+        "
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+fn many_args(a, b, c, d: d, e: a, f, g) {
+  todo
+}
+",
+        find_position_of("many_args").to_selection()
+    );
+}

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

@@ -0,0 +1,21 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "\nfn wibble(a: a, b: b, c: c) -> d { todo }\n\nfn many_args(a, b, c, d: d, e: a, f, g) {\n  todo\n}\n"
+---
+----- BEFORE ACTION
+
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+fn many_args(a, b, c, d: d, e: a, f, g) {
+   ↑                                     
+  todo
+}
+
+
+----- AFTER ACTION
+
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+fn many_args(a: b, b: c, c: e, d: d, e: a, f: f, g: g) -> h {
+  todo
+}

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

@@ -0,0 +1,21 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "\nfn wibble(a: a, b: b, c: c) -> d { todo }\n\nfn pair(a, b) {\n  #(a, b)\n}\n"
+---
+----- BEFORE ACTION
+
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+fn pair(a, b) {
+   ↑           
+  #(a, b)
+}
+
+
+----- AFTER ACTION
+
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+fn pair(a: a, b: b) -> #(a, b) {
+  #(a, b)
+}

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

@@ -0,0 +1,17 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "\nfn wibble(a: a, b: b, c: c) -> d { todo }\n\nconst empty = []\n"
+---
+----- BEFORE ACTION
+
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+const empty = []
+      ↑         
+
+
+----- AFTER ACTION
+
+fn wibble(a: a, b: b, c: c) -> d { todo }
+
+const empty: List(a) = []

compiler-core/src/type_/printer.rs

@@ -305,7 +305,7 @@ pub enum PrintMode {
 /// names that types and modules have been aliased with in the current module.
 #[derive(Debug)]
 pub struct Printer<'a> {
-    names: &'a Names,
+    pub names: &'a Names,
     uid: u64,
 
     /// Some type variables aren't bound to names, so when trying to print those,
@@ -333,6 +333,14 @@ impl<'a> Printer<'a> {
         }
     }
 
+    pub fn clear_type_variables(&mut self) {
+        self.printed_type_variable_names.clear();
+    }
+
+    pub fn register_type_variable(&mut self, name: EcoString) {
+        _ = self.printed_type_variable_names.insert(name);
+    }
+
     pub fn print_type(&mut self, type_: &Type) -> EcoString {
         let mut buffer = EcoString::new();
         self.print(type_, &mut buffer, PrintMode::Normal);