feat(ocamllsp): Code action "Remove type annotation"

CHANGES.md

@@ -7,6 +7,9 @@
   ([#1037](https://github.com/ocaml/ocaml-lsp/pull/1037)), fixes
   [#1036](https://github.com/ocaml/ocaml-lsp/issues/1036)
 
+## Features
+- Add "Remove type annotation" code action. (#1039)
+
 # 1.15.1
 
 ## Fixes

ocaml-lsp-server/src/code_actions.ml

@@ -84,6 +84,7 @@ let compute server (params : CodeActionParams.t) =
             [ Action_destruct.t state
             ; Action_inferred_intf.t state
             ; Action_type_annotate.t
+            ; Action_remove_type_annotation.t
             ; Action_construct.t
             ; Action_refactor_open.unqualify
             ; Action_refactor_open.qualify

ocaml-lsp-server/src/code_actions/action_remove_type_annotation.ml

@@ -0,0 +1,77 @@
+open Import
+
+let action_kind = "remove type annotation"
+
+let check_typeable_context pipeline pos_start =
+  let pos_start = Mpipeline.get_lexing_pos pipeline pos_start in
+  let typer = Mpipeline.typer_result pipeline in
+  let browse = Mbrowse.of_typedtree (Mtyper.get_typedtree typer) in
+  let is_exp_constrained = function
+    | Typedtree.Texp_constraint _, loc, _ -> Some loc
+    | Typedtree.Texp_coerce (Some { ctyp_loc; _ }, _), _, _ -> Some ctyp_loc
+    | _ -> None
+  in
+  let is_pat_constrained = function
+    | Typedtree.Tpat_constraint _, loc, _ -> Some loc
+    | _ -> None
+  in
+  let is_valid loc p extras =
+    match extras |> List.rev |> List.find_map ~f:p with
+    | Some x -> `Valid (loc, x)
+    | None -> `Invalid
+  in
+  match Mbrowse.enclosing pos_start [ browse ] with
+  | (_, Expression e) :: _ -> is_valid e.exp_loc is_exp_constrained e.exp_extra
+  | (_, Pattern { pat_desc = Typedtree.Tpat_any; pat_loc; _ })
+    :: (_, Pattern { pat_desc = Typedtree.Tpat_alias _; pat_extra; _ })
+    :: _ -> is_valid pat_loc is_pat_constrained pat_extra
+  | (_, Pattern p) :: _ -> is_valid p.pat_loc is_pat_constrained p.pat_extra
+  | _ :: _ | [] -> `Invalid
+
+let get_source_text doc (loc : Loc.t) =
+  let open Option.O in
+  let source = Document.source doc in
+  let* start = Position.of_lexical_position loc.loc_start in
+  let+ end_ = Position.of_lexical_position loc.loc_end in
+  let (`Offset start) = Msource.get_offset source (Position.logical start) in
+  let (`Offset end_) = Msource.get_offset source (Position.logical end_) in
+  String.sub (Msource.text source) ~pos:start ~len:(end_ - start)
+
+let code_action_of_type_enclosing uri doc (loc, constr_loc) =
+  let open Option.O in
+  let+ src_text = get_source_text doc loc in
+  let edit : WorkspaceEdit.t =
+    let textedit : TextEdit.t =
+      { range = Range.of_loc (Loc.union loc constr_loc); newText = src_text }
+    in
+    let version = Document.version doc in
+    let textDocument =
+      OptionalVersionedTextDocumentIdentifier.create ~uri ~version ()
+    in
+    let edit =
+      TextDocumentEdit.create ~textDocument ~edits:[ `TextEdit textedit ]
+    in
+    WorkspaceEdit.create ~documentChanges:[ `TextDocumentEdit edit ] ()
+  in
+  let title = String.capitalize_ascii action_kind in
+  CodeAction.create
+    ~title
+    ~kind:(CodeActionKind.Other action_kind)
+    ~edit
+    ~isPreferred:false
+    ()
+
+let code_action doc (params : CodeActionParams.t) =
+  match Document.kind doc with
+  | `Other -> Fiber.return None
+  | `Merlin merlin ->
+    let pos_start = Position.logical params.range.start in
+    Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
+        let context = check_typeable_context pipeline pos_start in
+        match context with
+        | `Invalid -> None
+        | `Valid (loc1, loc2) ->
+          code_action_of_type_enclosing params.textDocument.uri doc (loc1, loc2))
+
+let t =
+  { Code_action.kind = CodeActionKind.Other action_kind; run = code_action }

ocaml-lsp-server/src/code_actions/action_remove_type_annotation.mli

@@ -0,0 +1 @@
+val t : Code_action.t

ocaml-lsp-server/src/ocaml_lsp_server.ml

@@ -35,6 +35,7 @@ let initialize_info (client_capabilities : ClientCapabilities.t) :
       :: List.map
            ~f:(fun (c : Code_action.t) -> c.kind)
            [ Action_type_annotate.t
+           ; Action_remove_type_annotation.t
            ; Action_construct.t
            ; Action_refactor_open.unqualify
            ; Action_refactor_open.qualify

ocaml-lsp-server/test/e2e-new/code_actions.ml

@@ -70,6 +70,12 @@ let find_annotate_action =
   | `CodeAction { kind = Some (Other "type-annotate"); _ } -> true
   | _ -> false
 
+let find_remove_annotation_action =
+  let open CodeAction in
+  function
+  | `CodeAction { kind = Some (Other "remove type annotation"); _ } -> true
+  | _ -> false
+
 let%expect_test "code actions" =
   let source = {ocaml|
 let foo = 123
@@ -153,3 +159,152 @@ let f x = (1 : int :> int)
   in
   print_code_actions source range ~filter:find_annotate_action;
   [%expect {| No code actions |}]
+
+let%expect_test "can remove type annotation from a function argument" =
+  let source =
+    {ocaml|
+type t = Foo of int | Bar of bool
+let f (x : t) = Foo x
+|ocaml}
+  in
+  let range =
+    let start = Position.create ~line:2 ~character:7 in
+    let end_ = Position.create ~line:2 ~character:8 in
+    Range.create ~start ~end_
+  in
+  print_code_actions source range ~filter:find_remove_annotation_action;
+  [%expect
+    {|
+    Code actions:
+    {
+      "edit": {
+        "documentChanges": [
+          {
+            "edits": [
+              {
+                "newText": "x",
+                "range": {
+                  "end": { "character": 13, "line": 2 },
+                  "start": { "character": 6, "line": 2 }
+                }
+              }
+            ],
+            "textDocument": { "uri": "file:///foo.ml", "version": 0 }
+          }
+        ]
+      },
+      "isPreferred": false,
+      "kind": "remove type annotation",
+      "title": "Remove type annotation"
+    } |}]
+
+let%expect_test "can remove type annotation from a toplevel value" =
+  let source = {ocaml|
+let (iiii : int) = 3 + 4
+|ocaml} in
+  let range =
+    let start = Position.create ~line:1 ~character:5 in
+    let end_ = Position.create ~line:1 ~character:6 in
+    Range.create ~start ~end_
+  in
+  print_code_actions source range ~filter:find_remove_annotation_action;
+  [%expect
+    {|
+    Code actions:
+    {
+      "edit": {
+        "documentChanges": [
+          {
+            "edits": [
+              {
+                "newText": "iiii",
+                "range": {
+                  "end": { "character": 16, "line": 1 },
+                  "start": { "character": 4, "line": 1 }
+                }
+              }
+            ],
+            "textDocument": { "uri": "file:///foo.ml", "version": 0 }
+          }
+        ]
+      },
+      "isPreferred": false,
+      "kind": "remove type annotation",
+      "title": "Remove type annotation"
+    } |}]
+
+let%expect_test "can remove type annotation from an argument in a function call"
+    =
+  let source =
+    {ocaml|
+let f (x : int) = x + 1
+ let () =
+   let i = 8 in
+   print_int (f i)
+|ocaml}
+  in
+  let range =
+    let start = Position.create ~line:1 ~character:7 in
+    let end_ = Position.create ~line:1 ~character:8 in
+    Range.create ~start ~end_
+  in
+  print_code_actions source range ~filter:find_remove_annotation_action;
+  [%expect
+    {|
+    Code actions:
+    {
+      "edit": {
+        "documentChanges": [
+          {
+            "edits": [
+              {
+                "newText": "x",
+                "range": {
+                  "end": { "character": 15, "line": 1 },
+                  "start": { "character": 6, "line": 1 }
+                }
+              }
+            ],
+            "textDocument": { "uri": "file:///foo.ml", "version": 0 }
+          }
+        ]
+      },
+      "isPreferred": false,
+      "kind": "remove type annotation",
+      "title": "Remove type annotation"
+    } |}]
+
+let%expect_test "can remove type annotation from a coerced expression" =
+  let source = {ocaml|
+let x = (7 : int :> int)
+|ocaml} in
+  let range =
+    let start = Position.create ~line:1 ~character:9 in
+    let end_ = Position.create ~line:1 ~character:10 in
+    Range.create ~start ~end_
+  in
+  print_code_actions source range ~filter:find_remove_annotation_action;
+  [%expect
+    {|
+    Code actions:
+    {
+      "edit": {
+        "documentChanges": [
+          {
+            "edits": [
+              {
+                "newText": "7",
+                "range": {
+                  "end": { "character": 16, "line": 1 },
+                  "start": { "character": 9, "line": 1 }
+                }
+              }
+            ],
+            "textDocument": { "uri": "file:///foo.ml", "version": 0 }
+          }
+        ]
+      },
+      "isPreferred": false,
+      "kind": "remove type annotation",
+      "title": "Remove type annotation"
+    } |}]

ocaml-lsp-server/test/e2e-new/start_stop.ml

@@ -48,7 +48,8 @@ let%expect_test "start/stop" =
             "codeActionKinds": [
               "quickfix", "refactor.inline", "construct", "destruct",
               "inferred_intf", "put module name in identifiers",
-              "remove module name from identifiers", "type-annotate"
+              "remove module name from identifiers", "remove type annotation",
+              "type-annotate"
             ]
           },
           "codeLensProvider": { "resolveProvider": false },