feat(lsp): Explicit type annotation code action (#2125)

compiler/src/language_server/code_action.re

@@ -0,0 +1,105 @@
+open Grain;
+open Compile;
+open Grain_parsing;
+open Grain_utils;
+open Grain_typed;
+open Grain_diagnostics;
+open Sourcetree;
+open Lsp_types;
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#definitionParams
+module RequestParams = {
+  [@deriving yojson({strict: false})]
+  type code_action_context = {diagnostics: list(Protocol.diagnostic)};
+
+  [@deriving yojson({strict: false})]
+  type t = {
+    [@key "textDocument"]
+    text_document: Protocol.text_document_identifier,
+    range: Protocol.range,
+    context: code_action_context,
+  };
+};
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#locationLink
+module ResponseResult = {
+  [@deriving yojson]
+  type code_action = {
+    title: string,
+    kind: string,
+    edit: Protocol.workspace_edit,
+  };
+
+  [@deriving yojson]
+  type t = option(list(code_action));
+};
+
+let send_no_result = (~id: Protocol.message_id) => {
+  Protocol.response(~id, `Null);
+};
+
+let explicit_type_annotation = (range, uri, type_str) => {
+  ResponseResult.{
+    title: "Annotate type",
+    kind: "annotate-type",
+    edit: {
+      document_changes: [
+        {
+          text_document: {
+            uri,
+            version: None,
+          },
+          edits: [{range, new_text: ": " ++ type_str}],
+        },
+      ],
+    },
+  };
+};
+
+let send_code_actions =
+    (id: Protocol.message_id, code_actions: list(ResponseResult.code_action)) => {
+  Protocol.response(~id, ResponseResult.to_yojson(Some(code_actions)));
+};
+
+let process_explicit_type_annotation = (uri, results: list(Sourcetree.node)) => {
+  switch (results) {
+  | [Pattern({pattern})]
+  | [
+      Pattern({pattern: {pat_desc: TPatAlias({pat_desc: TPatAny}, _, _)}}),
+      Pattern({pattern}),
+      ..._,
+    ]
+      when pattern.pat_extra == [] =>
+    let loc = {...pattern.pat_loc, loc_start: pattern.pat_loc.loc_end};
+    let type_str = Printtyp.string_of_type_scheme(pattern.pat_type);
+    Some(explicit_type_annotation(Utils.loc_to_range(loc), uri, type_str));
+  | _ => None
+  };
+};
+
+let process =
+    (
+      ~id: Protocol.message_id,
+      ~compiled_code: Hashtbl.t(Protocol.uri, Lsp_types.code),
+      ~documents: Hashtbl.t(Protocol.uri, string),
+      params: RequestParams.t,
+    ) => {
+  switch (Hashtbl.find_opt(compiled_code, params.text_document.uri)) {
+  | None => send_no_result(~id)
+  | Some({program, sourcetree}) =>
+    let results = Sourcetree.query(params.range.range_start, sourcetree);
+
+    let code_actions =
+      List.filter_map(
+        x => x,
+        [
+          process_explicit_type_annotation(params.text_document.uri, results),
+        ],
+      );
+
+    switch (code_actions) {
+    | [] => send_no_result(~id)
+    | _ => send_code_actions(id, code_actions)
+    };
+  };
+};

compiler/src/language_server/code_action.rei

@@ -0,0 +1,22 @@
+open Grain_typed;
+
+//https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#documentFormattingParams
+module RequestParams: {
+  [@deriving yojson({strict: false})]
+  type t;
+};
+
+//https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textEdit
+module ResponseResult: {
+  [@deriving yojson]
+  type t;
+};
+
+let process:
+  (
+    ~id: Protocol.message_id,
+    ~compiled_code: Hashtbl.t(Protocol.uri, Lsp_types.code),
+    ~documents: Hashtbl.t(Protocol.uri, string),
+    RequestParams.t
+  ) =>
+  unit;

compiler/src/language_server/driver.re

@@ -49,6 +49,9 @@ let process = msg => {
   | Goto(id, goto_request_type, params) when is_initialized^ =>
     Goto.process(~id, ~compiled_code, ~documents, goto_request_type, params);
     Reading;
+  | CodeAction(id, params) when is_initialized^ =>
+    Code_action.process(~id, ~compiled_code, ~documents, params);
+    Reading;
   | SetTrace(trace_value) =>
     Trace.set_level(trace_value);
     Reading;

compiler/src/language_server/formatting.re

@@ -25,16 +25,9 @@ module RequestParams = {
   };
 };
 
-// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textEdit
 module ResponseResult = {
   [@deriving yojson]
-  type text_edit = {
-    range: Protocol.range,
-    newText: string,
-  };
-
-  [@deriving yojson]
-  type t = option(list(text_edit));
+  type t = option(list(Protocol.text_edit));
 };
 
 let process =
@@ -76,7 +69,8 @@ let process =
             },
         };
 
-        let res: ResponseResult.t = Some([{range, newText: formatted_code}]);
+        let res: ResponseResult.t =
+          Some([{range, new_text: formatted_code}]);
         Protocol.response(~id, ResponseResult.to_yojson(res));
       }) {
       | exn =>

compiler/src/language_server/initialize.re

@@ -75,7 +75,7 @@ module ResponseResult = {
     type_definition_provider: true,
     references_provider: false,
     document_symbol_provider: false,
-    code_action_provider: false,
+    code_action_provider: true,
     code_lens_provider: {
       resolve_provider: true,
     },

compiler/src/language_server/message.re

@@ -11,6 +11,7 @@ type t =
   | TextDocumentInlayHint(Protocol.message_id, Inlayhint.RequestParams.t)
   | Formatting(Protocol.message_id, Formatting.RequestParams.t)
   | Goto(Protocol.message_id, Goto.goto_request_type, Goto.RequestParams.t)
+  | CodeAction(Protocol.message_id, Code_action.RequestParams.t)
   | SetTrace(Protocol.trace_value)
   | Unsupported
   | Error(string);
@@ -76,6 +77,11 @@ let of_request = (msg: Protocol.request_message): t => {
     | Ok(params) => Goto(id, TypeDefinition, params)
     | Error(msg) => Error(msg)
     }
+  | {method: "textDocument/codeAction", id: Some(id), params: Some(params)} =>
+    switch (Code_action.RequestParams.of_yojson(params)) {
+    | Ok(params) => CodeAction(id, params)
+    | Error(msg) => Error(msg)
+    }
   | {method: "$/setTrace", params: Some(params)} =>
     switch (Trace.RequestParams.of_yojson(params)) {
     | Ok(params) => SetTrace(params.value)

compiler/src/language_server/message.rei

@@ -9,6 +9,7 @@ type t =
   | TextDocumentInlayHint(Protocol.message_id, Inlayhint.RequestParams.t)
   | Formatting(Protocol.message_id, Formatting.RequestParams.t)
   | Goto(Protocol.message_id, Goto.goto_request_type, Goto.RequestParams.t)
+  | CodeAction(Protocol.message_id, Code_action.RequestParams.t)
   | SetTrace(Protocol.trace_value)
   | Unsupported
   | Error(string);

compiler/src/language_server/protocol.re

@@ -116,6 +116,13 @@ type text_document_position_params = {
   position,
 };
 
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#optionalVersionedTextDocumentIdentifier
+[@deriving yojson]
+type optional_versioned_text_document_identifier = {
+  uri,
+  version: option(int),
+};
+
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocumentSyncKind
 [@deriving (enum, yojson)]
 type text_document_sync_kind =
@@ -216,6 +223,29 @@ type notification_message = {
   params: Yojson.Safe.t,
 };
 
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textEdit
+[@deriving yojson({strict: false})]
+type text_edit = {
+  range,
+  [@key "newText"]
+  new_text: string,
+};
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocumentEdit
+[@deriving yojson({strict: false})]
+type text_document_edit = {
+  [@key "textDocument"]
+  text_document: optional_versioned_text_document_identifier,
+  edits: list(text_edit),
+};
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspaceEdit
+[@deriving yojson({strict: false})]
+type workspace_edit = {
+  [@key "documentChanges"]
+  document_changes: list(text_document_edit),
+};
+
 let version: version = "2.0";
 
 let header_prefix = "Content-Length: ";

compiler/src/language_server/protocol.rei

@@ -122,6 +122,13 @@ type versioned_text_document_identifier = {
   version: int,
 };
 
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#optionalVersionedTextDocumentIdentifier
+[@deriving yojson({strict: false})]
+type optional_versioned_text_document_identifier = {
+  uri,
+  version: option(int),
+};
+
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#errorCodes
 [@deriving (enum, yojson)]
 type error_code =
@@ -170,6 +177,29 @@ type notification_message = {
   params: Yojson.Safe.t,
 };
 
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textEdit
+[@deriving yojson({strict: false})]
+type text_edit = {
+  range,
+  [@key "newText"]
+  new_text: string,
+};
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocumentEdit
+[@deriving yojson({strict: false})]
+type text_document_edit = {
+  [@key "textDocument"]
+  text_document: optional_versioned_text_document_identifier,
+  edits: list(text_edit),
+};
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspaceEdit
+[@deriving yojson({strict: false})]
+type workspace_edit = {
+  [@key "documentChanges"]
+  document_changes: list(text_document_edit),
+};
+
 //https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#definitionClientCapabilities
 [@deriving yojson({strict: false})]
 type definition_client_capabilities = {