Add wrap-type-in-module code action

ocaml-lsp-server/src/code_actions.ml

@@ -51,6 +51,7 @@ let compute_ocaml_code_actions (params : CodeActionParams.t) state doc =
       ; Action_inline.t
       ; Action_extract.local
       ; Action_extract.function_
+      ; Action_wrap_type_in_module.t
       ]
   in
   let batchable, non_batchable =

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

@@ -28,19 +28,9 @@ let check_typeable_context pipeline pos_start =
   | _ :: _ | [] -> `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, typ) =
   let open Option.O in
-  let+ original_text = get_source_text doc loc in
+  let+ original_text = Document.get_source_text doc loc in
   let newText = Printf.sprintf "(%s : %s)" original_text typ in
   let edit : WorkspaceEdit.t =
     let textedit : TextEdit.t = { range = Range.of_loc loc; newText } in

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

@@ -0,0 +1,119 @@
+open Import
+
+let action_kind = "wrap-type-in-module"
+
+(** Gets the type definition surrounding the cursor position if the cursor is within a
+    type definition. *)
+let type_definition_at 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 scopes = Mbrowse.enclosing pos_start [ browse ] in
+  List.find_map scopes ~f:(function
+    | _, Type_declaration type_decl -> Some type_decl
+    | _ -> None)
+;;
+
+(* Gets the portions of an outer interval that surround an inner subinterval. In an
+   analogous case with just character positions, surrounding_portions ~outer:[0:5]
+   ~inner:[2:3] = ([0, 1], [4, 5]).
+
+   [inner] must be strictly contained within [outer]. This will always be the case with
+   the type name of a type declaration. *)
+let surrounding_portions
+      ~(inner : Merlin_parsing.Location.t)
+      ~(outer : Merlin_parsing.Location.t)
+  =
+  let before : Merlin_parsing.Location.t =
+    { loc_start = outer.loc_start
+    ; loc_end = inner.loc_start
+    ; loc_ghost = outer.loc_ghost || inner.loc_ghost
+    }
+  in
+  let after : Merlin_parsing.Location.t =
+    { loc_start = inner.loc_end
+    ; loc_end = outer.loc_end
+    ; loc_ghost = outer.loc_ghost || inner.loc_ghost
+    }
+  in
+  before, after
+;;
+
+let leading_whitespace doc (loc : Loc.t) =
+  let before : Loc.t =
+    { loc_start = { loc.loc_start with pos_cnum = 0 }
+    ; loc_end = loc.loc_start
+    ; loc_ghost = false
+    }
+  in
+  Document.get_source_text doc before
+;;
+
+let new_module_text doc (type_decl : Ocaml_typing.Typedtree.type_declaration) =
+  let open Option.O in
+  let before, after =
+    surrounding_portions ~inner:type_decl.typ_name.loc ~outer:type_decl.typ_loc
+  in
+  let* before_text = Document.get_source_text doc before in
+  let* after_text = Document.get_source_text doc after in
+  let* original_indent = leading_whitespace doc type_decl.typ_loc in
+  let new_type_decl = before_text ^ "t" ^ after_text in
+  let indented_type_decl =
+    (* The type_decl is unevenly indented because the first line is unindented. This is
+       because the type_decl's location doesn't necessarily start at column 0. *)
+    original_indent ^ new_type_decl
+    |> String.split_lines
+    |> List.map ~f:(fun s -> "  " ^ s)
+    |> String.concat ~sep:"\n"
+  in
+  let module_name = String.capitalize_ascii type_decl.typ_name.txt in
+  match Document.kind doc with
+  | `Merlin m ->
+    let assign =
+      match Document.Merlin.kind m with
+      | Document.Kind.Intf -> ": sig"
+      | Document.Kind.Impl -> "= struct"
+    in
+    Some
+      (String.concat
+         ~sep:"\n"
+         [ (* Don't indent the first line because the edit starts at the original "t" in
+              "type". *)
+           Printf.sprintf "module %s %s" module_name assign
+         ; indented_type_decl
+         ; original_indent ^ "end"
+         ])
+  | `Other -> None
+;;
+
+let code_action pipeline doc (params : CodeActionParams.t) =
+  let open Option.O in
+  let pos_start = Position.logical params.range.start in
+  let* type_decl = type_definition_at pipeline pos_start in
+  let type_name = type_decl.typ_name.txt in
+  if String.equal type_name "t"
+  then None
+  else
+    let* newText = new_module_text doc type_decl in
+    let uri = params.textDocument.uri in
+    let edit : WorkspaceEdit.t =
+      let textedit : TextEdit.t = { range = Range.of_loc type_decl.typ_loc; newText } 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
+    Some
+      (CodeAction.create
+         ~title
+         ~kind:(CodeActionKind.Other action_kind)
+         ~edit
+         ~isPreferred:false
+         ())
+;;
+
+let kind = CodeActionKind.Other action_kind
+let t = Code_action.batchable kind code_action

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

@@ -0,0 +1,24 @@
+(** Code action to wrap a type in a module. Useful for not having to write boilerplate
+    module ... = struct ... end every time you define a type. Example:
+    {[
+      type 'a foo =
+        { a : int
+        ; b : int
+        ; c : 'a list
+        }
+      [@@deriving sexp]
+    ]}
+    becomes
+    {[
+      module Foo = struct
+        type 'a t =
+          { a : int
+          ; b : int
+          ; c : 'a list
+          }
+        [@@deriving sexp]
+      end
+    ]}
+    after the action. Types already named "t" are ignored. *)
+
+val t : Code_action.t

ocaml-lsp-server/src/document.ml

@@ -471,3 +471,13 @@ let substring doc range =
   then None
   else Some (String.sub text ~pos:start ~len:(end_ - start))
 ;;
+
+let get_source_text doc (loc : Loc.t) =
+  let open Option.O in
+  let source = 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)
+;;

ocaml-lsp-server/src/document.mli

@@ -123,3 +123,9 @@ val edit : t -> TextEdit.t list -> WorkspaceEdit.t
 
     Returns [None] when there is no corresponding substring. *)
 val substring : t -> Range.t -> string option
+
+(** [get_source_text t loc] returns the substring of the document [t] that corresponds to
+    the location [loc].
+
+    Returns [None] when there is no corresponding substring. *)
+val get_source_text : t -> Loc.t -> string option

ocaml-lsp-server/src/testing.ml

@@ -4,3 +4,4 @@ module Compl = Compl
 module Merlin_kernel = Merlin_kernel
 module Prefix_parser = Prefix_parser
 module Range = Range
+module Action_wrap_type_in_module = Action_wrap_type_in_module

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

@@ -0,0 +1,178 @@
+let wrap_test = Code_actions.code_action_test ~title:"Wrap-type-in-module"
+
+let%expect_test "preserving whitespace" =
+  wrap_test "t$y$pe foo = bar";
+  [%expect
+    {|
+    module Foo = struct
+      type t = bar
+    end
+    |}];
+  wrap_test "t$y$pe foo= bar";
+  [%expect
+    {|
+    module Foo = struct
+      type t= bar
+    end
+    |}];
+  wrap_test "t$y$pe foo\n=\nbar";
+  [%expect
+    {|
+    module Foo = struct
+      type t
+      =
+      bar
+    end
+    |}];
+  wrap_test
+      (
+    String.concat
+      "\n"
+      [ "typ$e$ ('a, 'b) a = { bar : 'a";
+        "                  ; baz : 'b";
+        "                  }" ]
+      [@ocamlformat "disable"]);
+  [%expect
+    {|
+    module A = struct
+      type ('a, 'b) t = { bar : 'a
+                        ; baz : 'b
+                        }
+    end
+    |}];
+  (* non-space character can come before the type name *)
+  wrap_test
+    (String.concat
+      "\n"
+      [ "ty$p$e";
+        "abc = { a: int;";
+        "b: int; c: int } [@@deriving sexp]" ]) [@ocamlformat "disable"];
+  [%expect
+    {|
+    module Abc = struct
+      type
+      t = { a: int;
+      b: int; c: int } [@@deriving sexp]
+    end
+    |}];
+  (* entire module is indented by correct amount *)
+  wrap_test
+    (String.concat
+        "\n"
+        [ "module Outer = struct"
+        ; "  module Inner = struct"
+        ; "    type record ="
+        ; "      { foo : int"
+        ; "      ; bar : in$t$"
+        ; "      }"
+        ; "  end"
+        ; "end"
+        ]);
+  [%expect
+    {|
+    module Outer = struct
+      module Inner = struct
+        module Record = struct
+          type t =
+            { foo : int
+            ; bar : int
+            }
+        end
+      end
+    end
+    |}];
+;;
+
+let%expect_test "type parameters" =
+  wrap_test "t$y$pe ('a, 'b, _) foo = bar";
+  [%expect
+    {|
+    module Foo = struct
+      type ('a, 'b, _) t = bar
+    end
+    |}];
+;;
+
+let%expect_test "definition chain" =
+    wrap_test
+      {xxx|module Module = struct
+  module Foo = struct
+    type ('a, 'b) t =
+      { a : 'a
+      ; b : 'b
+      }
+  end
+end
+
+ty$p$e ('a, 'b) foo = ('a, 'b) Module.Foo.t =
+  { a : 'a
+  ; b : 'b
+  }
+|xxx};
+  [%expect
+    {|
+    module Module = struct
+      module Foo = struct
+        type ('a, 'b) t =
+          { a : 'a
+          ; b : 'b
+          }
+      end
+    end
+
+    module Foo = struct
+      type ('a, 'b) t = ('a, 'b) Module.Foo.t =
+        { a : 'a
+        ; b : 'b
+        }
+    end
+    |}];
+;;
+
+let%expect_test "can trigger action on any part of type declaration" =
+  wrap_test {|type abc = { a: int; b: int; c: int } [@@derivin$g$ sexp]|};
+  [%expect
+    {|
+    module Abc = struct
+      type t = { a: int; b: int; c: int } [@@deriving sexp]
+    end
+    |}];
+  wrap_test {|type abc = { a: int; b: int;$ $c: int } [@@deriving sexp]|};
+  [%expect
+    {|
+    module Abc = struct
+      type t = { a: int; b: int; c: int } [@@deriving sexp]
+    end
+    |}];
+  wrap_test {|type a$b$c = { a: int; b: int; c: int } [@@deriving sexp]|};
+  [%expect
+    {|
+    module Abc = struct
+      type t = { a: int; b: int; c: int } [@@deriving sexp]
+    end
+    |}];
+  wrap_test {|typ$e$ abc = { a: int; b: int; c: int } [@@deriving sexp]|};
+  [%expect
+    {|
+    module Abc = struct
+      type t = { a: int; b: int; c: int } [@@deriving sexp]
+    end
+    |}];
+;;
+
+let%expect_test "type with name t is ignored" =
+  wrap_test {|type $t$ = int|};
+  [%expect {| |}]
+;;
+
+let%expect_test "produce sig in mli" =
+  wrap_test
+    ?path:(Some "needs-refactoring.mli")
+    {|type abc = { a: int; b: int; c: int } [@@derivin$g$ sexp]|};
+  [%expect
+    {|
+    module Abc : sig
+      type t = { a: int; b: int; c: int } [@@deriving sexp]
+    end
+    |}]
+;;