feat(LS): Annotate all top level definitions code action

CHANGELOG.md

@@ -2,6 +2,43 @@
 
 ## Unreleased
 
+### Language server
+
+- The language server now offers code action to add type annotations to all
+  functions and constants. For example,
+
+  ```gleam
+  pub const answer = 42
+
+  pub fn add(x, y) {
+    x + y
+  }
+
+  pub fn add_one(thing) {
+    //     ^ Triggering "Annotate all top level definitions" code action here
+    let result = add(thing, 1)
+    result
+  }
+  ```
+
+  Triggering the "Annotate all top level definitions" code action over
+  the name of function `add_one` would result in following code:
+
+  ```gleam
+  pub const answer: Int = 42
+
+  pub fn add(x: Int, y: Int) -> Int {
+    x + y
+  }
+
+  pub fn add_one(thing: Int) -> Int {
+    let result = add(thing, 1)
+    result
+  }
+  ```
+
+  ([Andrey Kozhev](https://github.com/ankddev))
+
 ### Bug fixes
 
 - Fixed a bug where the "Extract function" code action would not properly

compiler-core/src/language_server/code_action.rs

@@ -1360,6 +1360,113 @@ impl<'a> AddAnnotations<'a> {
     }
 }
 
+/// Code action to add type annotations to all top level definitions
+///
+pub struct AnnotateTopLevelDefinitions<'a> {
+    module: &'a Module,
+    params: &'a CodeActionParams,
+    edits: TextEdits<'a>,
+    is_hovering_definition: bool,
+}
+
+impl<'a> AnnotateTopLevelDefinitions<'a> {
+    pub fn new(
+        module: &'a Module,
+        line_numbers: &'a LineNumbers,
+        params: &'a CodeActionParams,
+    ) -> Self {
+        Self {
+            module,
+            params,
+            edits: TextEdits::new(line_numbers),
+            is_hovering_definition: false,
+        }
+    }
+
+    pub fn code_actions(mut self) -> Vec<CodeAction> {
+        self.visit_typed_module(&self.module.ast);
+
+        // We only want to trigger the action if we're over one of the definition in
+        // the module
+        if !self.is_hovering_definition || self.edits.edits.is_empty() {
+            return vec![];
+        };
+
+        let mut action = Vec::with_capacity(1);
+        CodeActionBuilder::new("Annotate all top level definitions")
+            .kind(CodeActionKind::REFACTOR_REWRITE)
+            .changes(self.params.text_document.uri.clone(), self.edits.edits)
+            .preferred(false)
+            .push_to(&mut action);
+        action
+    }
+}
+
+impl<'ast> ast::visit::Visit<'ast> for AnnotateTopLevelDefinitions<'_> {
+    fn visit_typed_module_constant(&mut self, constant: &'ast TypedModuleConstant) {
+        let code_action_range = self.edits.src_span_to_lsp_range(constant.location);
+
+        if overlaps(code_action_range, self.params.range) {
+            self.is_hovering_definition = true;
+        }
+
+        // We don't need to add an annotation if there already is one
+        if constant.annotation.is_some() {
+            return;
+        }
+
+        self.edits.insert(
+            constant.name_location.end,
+            format!(
+                ": {}",
+                // Create new printer to ignore type variables from other definitions
+                Printer::new_without_type_variables(&self.module.ast.names)
+                    .print_type(&constant.type_)
+            ),
+        );
+    }
+
+    fn visit_typed_function(&mut self, fun: &'ast ast::TypedFunction) {
+        // Create new printer to ignore type variables from other definitions
+        let mut printer = Printer::new_without_type_variables(&self.module.ast.names);
+        collect_type_variables(&mut printer, fun);
+
+        let code_action_range = self.edits.src_span_to_lsp_range(
+            fun.body_start
+                .map(|body_start| SrcSpan {
+                    start: fun.location.start,
+                    end: body_start,
+                })
+                .unwrap_or(fun.location),
+        );
+
+        if overlaps(code_action_range, self.params.range) {
+            self.is_hovering_definition = true;
+        }
+
+        // Annotate each argument separately
+        for argument in fun.arguments.iter() {
+            // Don't annotate the argument if it's already annotated
+            if argument.annotation.is_some() {
+                continue;
+            }
+
+            self.edits.insert(
+                argument.location.end,
+                format!(": {}", printer.print_type(&argument.type_)),
+            );
+        }
+
+        // Annotate the return type if it isn't already annotated
+        if fun.return_annotation.is_none() {
+            self.edits.insert(
+                fun.location.end,
+                format!(" -> {}", printer.print_type(&fun.return_type)),
+            );
+        }
+    }
+}
+
 struct TypeVariableCollector<'a, 'b> {
     printer: &'a mut Printer<'b>,
 }

compiler-core/src/language_server/engine.rs

@@ -44,9 +44,9 @@ use std::{collections::HashSet, sync::Arc};
 use super::{
     DownloadDependencies, MakeLocker,
     code_action::{
-        AddAnnotations, CodeActionBuilder, ConvertFromUse, ConvertToFunctionCall, ConvertToPipe,
-        ConvertToUse, ExpandFunctionCapture, ExtractConstant, ExtractVariable,
-        FillInMissingLabelledArgs, FillUnusedFields, FixBinaryOperation,
+        AddAnnotations, AnnotateTopLevelDefinitions, CodeActionBuilder, ConvertFromUse,
+        ConvertToFunctionCall, ConvertToPipe, ConvertToUse, ExpandFunctionCapture, ExtractConstant,
+        ExtractVariable, FillInMissingLabelledArgs, FillUnusedFields, FixBinaryOperation,
         FixTruncatedBitArraySegment, GenerateDynamicDecoder, GenerateFunction, GenerateJsonEncoder,
         GenerateVariant, InlineVariable, InterpolateString, LetAssertToCase, PatternMatchOnValue,
         RedundantTupleInCaseSubject, RemoveEchos, RemoveUnusedImports, UseLabelShorthandSyntax,
@@ -461,6 +461,8 @@ where
             )
             .code_actions();
             AddAnnotations::new(module, &lines, &params).code_action(&mut actions);
+            actions
+                .extend(AnnotateTopLevelDefinitions::new(module, &lines, &params).code_actions());
             Ok(if actions.is_empty() {
                 None
             } else {

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

@@ -112,6 +112,7 @@ const ASSIGN_UNUSED_RESULT: &str = "Assign unused Result value to `_`";
 const ADD_MISSING_PATTERNS: &str = "Add missing patterns";
 const ADD_ANNOTATION: &str = "Add type annotation";
 const ADD_ANNOTATIONS: &str = "Add type annotations";
+const ANNOTATE_TOP_LEVEL_DEFINITIONS: &str = "Annotate all top level definitions";
 const CONVERT_FROM_USE: &str = "Convert from `use`";
 const CONVERT_TO_USE: &str = "Convert to `use`";
 const EXTRACT_VARIABLE: &str = "Extract variable";
@@ -10686,3 +10687,132 @@ pub fn main() -> Nil {
         find_position_of("function").to_selection()
     );
 }
+
+#[test]
+fn annotate_all_top_level_definitions_constant() {
+    assert_code_action!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+pub const answer = 42
+
+pub fn add_two(thing) {
+  thing + 2
+}
+
+pub fn add_one(thing) {
+  thing + 1
+}
+"#,
+        find_position_of("const").select_until(find_position_of("="))
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_function() {
+    assert_code_action!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+pub fn add_two(thing) {
+  thing + 2
+}
+
+pub fn add_one(thing) {
+  thing + 1
+}
+"#,
+        find_position_of("fn").select_until(find_position_of("("))
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_already_annotated() {
+    assert_no_code_actions!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+pub const answer: Int = 42
+
+pub fn add_two(thing: Int) -> Int {
+  thing + 2
+}
+
+pub fn add_one(thing: Int) -> Int {
+  thing + 1
+}
+"#,
+        find_position_of("fn").select_until(find_position_of("("))
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_inside_body() {
+    assert_no_code_actions!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+pub fn add_one(thing) {
+  thing + 1
+}
+"#,
+        find_position_of("thing + 1").to_selection()
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_partially_annotated() {
+    assert_code_action!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+pub const answer: Int = 42
+pub const another_answer = 43
+
+pub fn add_two(thing) -> Int {
+  thing + 2
+}
+
+pub fn add_one(thing: Int) {
+  thing + 1
+}
+"#,
+        find_position_of("fn").select_until(find_position_of("("))
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_with_partially_annotated_generic_function() {
+    assert_code_action!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+pub fn wibble(a: a, b, c: c, d) {
+  todo
+}
+"#,
+        find_position_of("wibble").to_selection()
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_with_two_generic_functions() {
+    assert_code_action!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+fn wibble(one) { todo }
+
+fn wobble(other) { todo }
+"#,
+        find_position_of("wobble").to_selection()
+    );
+}
+
+#[test]
+fn annotate_all_top_level_definitions_with_constant_and_generic_functions() {
+    assert_code_action!(
+        ANNOTATE_TOP_LEVEL_DEFINITIONS,
+        r#"
+const answer = 42
+
+fn wibble(one) { todo }
+
+fn wobble(other) { todo }
+"#,
+        find_position_of("wobble").to_selection()
+    );
+}

compiler-core/src/language_server/tests/snapshots/gleam_core__language_server__tests__action__add_multiple_annotations.snap.new

@@ -0,0 +1,29 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+assertion_line: 2928
+expression: "\npub const my_constant = 20\n\npub fn add_my_constant(value) {\n  let result = value + my_constant\n  result\n}\n"
+snapshot_kind: text
+---
+----- BEFORE ACTION
+
+pub const my_constant = 20
+▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔
+
+pub fn add_my_constant(value) {
+▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔
+  let result = value + my_constant
+▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔
+  result
+▔▔▔▔▔▔▔▔
+}
+↑
+
+
+----- AFTER ACTION
+
+pub const my_constant: Int = 20
+
+pub fn add_my_constant(value: Int) -> Int {
+  let result = value + my_constant
+  result
+}

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

@@ -0,0 +1,29 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "\npub const answer = 42\n\npub fn add_two(thing) {\n  thing + 2\n}\n\npub fn add_one(thing) {\n  thing + 1\n}\n"
+---
+----- BEFORE ACTION
+
+pub const answer = 42
+    ▔▔▔▔▔▔▔▔▔▔▔▔▔↑   
+
+pub fn add_two(thing) {
+  thing + 2
+}
+
+pub fn add_one(thing) {
+  thing + 1
+}
+
+
+----- AFTER ACTION
+
+pub const answer: Int = 42
+
+pub fn add_two(thing: Int) -> Int {
+  thing + 2
+}
+
+pub fn add_one(thing: Int) -> Int {
+  thing + 1
+}