Add code action to annotate declarations

crates/language_server/Cargo.toml

@@ -25,6 +25,7 @@ roc_solve_problem = { path = "../compiler/solve_problem" }
 roc_target = { path = "../compiler/roc_target" }
 roc_types = { path = "../compiler/types" }
 roc_packaging = {path = "../packaging"}
+roc_cli = { path = "../cli" }
 
 bumpalo.workspace = true
 parking_lot.workspace = true

crates/language_server/src/analysis/analysed_doc.rs

@@ -1,20 +1,28 @@
 use log::{debug, info};
+use roc_can::{
+    def::Def,
+    expr::Expr,
+    traverse::{find_declaration_at, DeclarationInfo, FoundDeclaration},
+};
+use roc_cli::{annotation_edit, annotation_edits};
+use roc_types::subs::Content;
 use std::collections::HashMap;
 
 use bumpalo::Bump;
 
 use roc_module::symbol::{ModuleId, Symbol};
 
-use roc_region::all::LineInfo;
+use roc_region::all::{LineColumn, LineInfo};
 
 use tower_lsp::lsp_types::{
-    CompletionItem, Diagnostic, GotoDefinitionResponse, Hover, HoverContents, LanguageString,
-    Location, MarkedString, Position, Range, SemanticTokens, SemanticTokensResult, TextEdit, Url,
+    CodeAction, CodeActionKind, CompletionItem, Diagnostic, GotoDefinitionResponse, Hover,
+    HoverContents, LanguageString, Location, MarkedString, Position, Range, SemanticTokens,
+    SemanticTokensResult, TextEdit, Url, WorkspaceEdit,
 };
 
 use crate::{
     analysis::completion::{field_completion, get_completion_items, get_module_completion_items},
-    convert::{ToRange, ToRocPosition},
+    convert::{ToRange, ToRegion, ToRocPosition},
 };
 
 use super::{
@@ -308,4 +316,120 @@ impl AnalyzedDocument {
             }
         }
     }
+
+    pub fn annotate(&self, range: Range) -> Option<CodeAction> {
+        let region = range.to_region(self.line_info());
+        let AnalyzedModule {
+            module_id,
+            interns,
+            subs,
+            abilities,
+            declarations,
+            ..
+        } = self.module()?;
+
+        if let Some(found_decl) = find_declaration_at(region, declarations) {
+            self.annotate_declaration(found_decl)
+        } else {
+            let edits = annotation_edits(
+                declarations,
+                subs,
+                abilities,
+                &self.doc_info.source,
+                *module_id,
+                interns,
+            )
+            .into_iter()
+            .map(|(offset, edit)| {
+                let LineColumn { line, column } = self.line_info().convert_offset(offset as u32);
+                let position = Position::new(line, column);
+                let range = Range::new(position, position);
+
+                TextEdit {
+                    range,
+                    new_text: edit,
+                }
+            })
+            .collect();
+
+            Some(CodeAction {
+                title: "Add top-level signatures".to_owned(),
+                edit: Some(WorkspaceEdit::new(HashMap::from([(
+                    self.url().clone(),
+                    edits,
+                )]))),
+                kind: Some(CodeActionKind::SOURCE),
+                ..Default::default()
+            })
+        }
+    }
+
+    fn annotate_declaration(&self, decl: FoundDeclaration<'_>) -> Option<CodeAction> {
+        let AnalyzedModule {
+            module_id,
+            interns,
+            subs,
+            ..
+        } = self.module()?;
+
+        use DeclarationInfo as DI;
+        use FoundDeclaration as FD;
+        let symbol_range =
+            match decl {
+                FD::Decl(DI::Value { loc_symbol, .. })
+                | FD::Decl(DI::Function { loc_symbol, .. }) => loc_symbol.byte_range(),
+                FD::Decl(DI::Destructure { loc_pattern, .. })
+                | FD::Def(Def { loc_pattern, .. }) => loc_pattern.byte_range(),
+                _ => return None,
+            };
+
+        let expr = match decl {
+            FD::Decl(DI::Value { loc_expr: expr, .. })
+            | FD::Decl(DI::Function { loc_body: expr, .. })
+            | FD::Decl(DI::Destructure { loc_expr: expr, .. })
+            | FD::Def(Def { loc_expr: expr, .. }) => &expr.value,
+            _ => return None,
+        };
+
+        let annotation = match decl {
+            FD::Decl(DI::Value { annotation, .. })
+            | FD::Decl(DI::Function { annotation, .. })
+            | FD::Decl(DI::Destructure { annotation, .. }) => annotation,
+            FD::Def(Def { annotation, .. }) => annotation.as_ref(),
+            _ => return None,
+        };
+
+        if annotation.is_some()
+            | matches!(expr, Expr::ImportParams(..))
+            | matches!(
+                subs.get_content_without_compacting(decl.var()),
+                Content::Error
+            )
+        {
+            return None;
+        }
+
+        let (offset, new_text) = annotation_edit(
+            &self.doc_info.source,
+            subs,
+            interns,
+            *module_id,
+            decl.var(),
+            symbol_range,
+        );
+
+        let LineColumn { line, column } = self.line_info().convert_offset(offset as u32);
+        let position = Position::new(line, column);
+        let range = Range::new(position, position);
+
+        let edit = TextEdit { range, new_text };
+        Some(CodeAction {
+            title: "Add signature".to_owned(),
+            edit: Some(WorkspaceEdit::new(HashMap::from([(
+                self.url().clone(),
+                vec![edit],
+            )]))),
+            ..Default::default()
+        })
+    }
 }

crates/language_server/src/registry.rs

@@ -9,8 +9,8 @@ use std::{
 use tokio::sync::{Mutex, MutexGuard};
 
 use tower_lsp::lsp_types::{
-    CompletionResponse, Diagnostic, GotoDefinitionResponse, Hover, Position, SemanticTokensResult,
-    TextEdit, Url,
+    CodeActionOrCommand, CodeActionResponse, CompletionResponse, Diagnostic,
+    GotoDefinitionResponse, Hover, Position, Range, SemanticTokensResult, TextEdit, Url,
 };
 
 use crate::analysis::{AnalyzedDocument, DocInfo};
@@ -217,4 +217,14 @@ impl Registry {
 
         Some(CompletionResponse::Array(completions))
     }
+
+    pub async fn code_actions(&self, url: &Url, range: Range) -> Option<CodeActionResponse> {
+        let document = self.latest_document_by_url(url).await?;
+
+        let mut responses = vec![];
+        if let Some(edit) = document.annotate(range) {
+            responses.push(CodeActionOrCommand::CodeAction(edit));
+        }
+        Some(responses)
+    }
 }

crates/language_server/src/server.rs

@@ -103,13 +103,15 @@ impl RocServer {
                 work_done_progress: None,
             },
         };
+        let code_action_provider = CodeActionProviderCapability::Simple(true);
         ServerCapabilities {
             text_document_sync: Some(text_document_sync),
             hover_provider: Some(hover_provider),
             definition_provider: Some(OneOf::Right(definition_provider)),
             document_formatting_provider: Some(OneOf::Right(document_formatting_provider)),
             semantic_tokens_provider: Some(semantic_tokens_provider),
             completion_provider: Some(completion_provider),
+            code_action_provider: Some(code_action_provider),
             ..ServerCapabilities::default()
         }
     }
@@ -332,6 +334,18 @@ impl LanguageServer for RocServer {
         )
         .await
     }
+
+    async fn code_action(&self, params: CodeActionParams) -> Result<Option<CodeActionResponse>> {
+        let CodeActionParams {
+            text_document,
+            range,
+            context: _,
+            partial_result_params: _,
+            work_done_progress_params: _,
+        } = params;
+
+        unwind_async(self.state.registry.code_actions(&text_document.uri, range)).await
+    }
 }
 
 async fn unwind_async<Fut, T>(future: Fut) -> tower_lsp::jsonrpc::Result<T>