Fix jump to hole on destruct/construct to jump only within the new text (#559)

* add `Document.new_range_on_replace`
* jump to a hole on destruct/construct only if
the hole is within the generated text
* make `Compl.complete` independent,
ie it takes state and completion params and returns completion response
* send jump to next hole only if the client has an experimental capability "jumpToNextHole"
* add support for experimental client capabilities
* jump to next hole for any client that has an experimental capability "jumpToNextHole"
* rename `CustomCommands` to `Custom_commands`

Author: ulugbekna

ocaml-lsp-server/src/client.ml

@@ -1,33 +1,45 @@
 open Import
 
+module Experimental_capabilities = struct
+  type t = bool
+
+  let of_opt_json (json : Json.t option) =
+    match json with
+    | Some (`Assoc fields) ->
+      Json.field fields "jumpToNextHole" Json.Conv.bool_of_yojson
+      |> Option.value ~default:false
+    | _ -> false
+
+  let supportsJumpToNextHole t = t
+end
+
 module Vscode = struct
   module Commands = struct
     let triggerSuggest =
       Command.create ~title:"Trigger Suggest"
         ~command:"editor.action.triggerSuggest" ()
+  end
+end
 
-    module Custom = struct
-      let next_hole ?start_position ~notify_if_no_hole () =
-        let arguments =
-          let arg_obj_fields =
-            let notif_json =
-              Some ("notify-if-no-hole", Json.bool notify_if_no_hole)
-            in
-            let pos_json =
-              Option.map start_position ~f:(fun p ->
-                  ("position", Position.yojson_of_t p))
-            in
-            List.filter_opt [ pos_json; notif_json ]
-          in
-          match arg_obj_fields with
-          | [] -> [] (* no arguments -- the extension uses defaults *)
-          | fields ->
-            (* the use of a (json) object as the first and single argument to
-               the command is intended *)
-            [ `Assoc fields ]
+module Custom_commands = struct
+  let next_hole ?in_range ~notify_if_no_hole () =
+    let arguments =
+      let arg_obj_fields =
+        let notif_json =
+          Some ("shouldNotifyIfNoHole", Json.bool notify_if_no_hole)
         in
-        Command.create ~title:"Jump to Next Hole" ~command:"ocaml.next-hole"
-          ~arguments ()
-    end
-  end
+        let in_range_json =
+          Option.map in_range ~f:(fun r -> ("inRange", Range.yojson_of_t r))
+        in
+        List.filter_opt [ in_range_json; notif_json ]
+      in
+      match arg_obj_fields with
+      | [] -> [] (* no arguments -- the extension uses defaults *)
+      | fields ->
+        (* the use of a (json) object as the first and single argument to the
+           command is intended *)
+        [ `Assoc fields ]
+    in
+    Command.create ~title:"Jump to Next Hole" ~command:"ocaml.next-hole"
+      ~arguments ()
 end

ocaml-lsp-server/src/client.mli

@@ -3,6 +3,16 @@ open Import
 (** This module is a collection of client-specific functionality (client =
     editor) *)
 
+module Experimental_capabilities : sig
+  (** Module to store experimental client capabilities *)
+
+  type t
+
+  val of_opt_json : Json.t option -> t
+
+  val supportsJumpToNextHole : t -> bool
+end
+
 module Vscode : sig
   (** A collection of VS Code editor commands.
 
@@ -12,24 +22,22 @@ module Vscode : sig
     (** [editor.action.triggerSuggest] is a vscode-specific command, which
         triggers the completion request on all completion providers *)
     val triggerSuggest : Command.t
+  end
+end
 
-    (** Represents custom commands, i.e., commands added by a certain extension.
-
-        Currently, the module includes custom commands introduced by "OCaml
-        Platform" extension *)
-    module Custom : sig
-      (** Request client cursor to jump to the next hole.
+(** Represents custom commands, i.e., commands added by a certain extension. *)
+module Custom_commands : sig
+  (** Request client cursor to jump to the next hole.
 
-          Looks for a hole starting at position [start_position], if provided;
-          otherwise, uses the cursor position.
+      See the documentation for this command in [vscode-ocaml-platform] for
+      details.
 
-          Will not show a pop-up notification if [notify-if-no-hole] is set to
-          [false] (the default value is [true]) *)
-      val next_hole :
-           ?start_position:Position.t
-        -> notify_if_no_hole:bool
-        -> unit
-        -> Command.t
-    end
-  end
+      @param in_range
+        to pick a hole only in a given range; if omitted, the whole document is
+        used
+      @param notify_if_no_hole
+        specifies whether we want the client to show the user a message if there
+        is no hole to jump to *)
+  val next_hole :
+    ?in_range:Range.t -> notify_if_no_hole:bool -> unit -> Command.t
 end

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

@@ -5,10 +5,11 @@ let action_kind = "destruct"
 
 let kind = CodeActionKind.Other action_kind
 
-let code_action_of_case_analysis doc uri (loc, newText) =
+let code_action_of_case_analysis ~supportsJumpToNextHole doc uri (loc, newText)
+    =
   let range : Range.t = Range.of_loc loc in
+  let textedit : TextEdit.t = { range; newText } in
   let edit : WorkspaceEdit.t =
-    let textedit : TextEdit.t = { range; newText } in
     let version = Document.version doc in
     let textDocument =
       OptionalVersionedTextDocumentIdentifier.create ~uri ~version ()
@@ -20,11 +21,16 @@ let code_action_of_case_analysis doc uri (loc, newText) =
   in
   let title = String.capitalize_ascii action_kind in
   let command =
-    Client.Vscode.Commands.Custom.next_hole ~start_position:range.start
-      ~notify_if_no_hole:false ()
+    if supportsJumpToNextHole then
+      Some
+        (Client.Custom_commands.next_hole
+           ~in_range:(Range.resize_for_edit textedit)
+           ~notify_if_no_hole:false ())
+    else
+      None
   in
   CodeAction.create ~title ~kind:(CodeActionKind.Other action_kind) ~edit
-    ~command ~isPreferred:false ()
+    ?command ~isPreferred:false ()
 
 let code_action (state : State.t) doc (params : CodeActionParams.t) =
   let uri = params.textDocument.uri in
@@ -47,7 +53,13 @@ let code_action (state : State.t) doc (params : CodeActionParams.t) =
         | Ok formatted_text -> formatted_text
         | Error _ -> newText
       in
-      Some (code_action_of_case_analysis doc uri (loc, newText))
+      let supportsJumpToNextHole =
+        State.experimental_client_capabilities state
+        |> Client.Experimental_capabilities.supportsJumpToNextHole
+      in
+      Some
+        (code_action_of_case_analysis ~supportsJumpToNextHole doc uri
+           (loc, newText))
     | Error
         { exn =
             ( Merlin_analysis.Destruct.Wrong_parent _ | Query_commands.No_nodes

ocaml-lsp-server/src/compl.ml

@@ -215,7 +215,7 @@ module Complete_with_construct = struct
       None
     | Error exn -> Exn_with_backtrace.reraise exn
 
-  let process_dispatch_resp = function
+  let process_dispatch_resp ~supportsJumpToNextHole = function
     | None -> []
     | Some (loc, constructed_exprs) ->
       let range = Range.of_loc loc in
@@ -231,19 +231,26 @@ module Complete_with_construct = struct
       in
       let completionItem_of_constructed_expr idx expr =
         let expr_wo_parens = deparen_constr_expr expr in
-        let textEdit = `TextEdit { TextEdit.range; newText = expr } in
+        let edit = { TextEdit.range; newText = expr } in
         let command =
-          Client.Vscode.Commands.Custom.next_hole ~start_position:range.start
-            ~notify_if_no_hole:false ()
+          if supportsJumpToNextHole then
+            Some
+              (Client.Custom_commands.next_hole
+                 ~in_range:(Range.resize_for_edit edit)
+                 ~notify_if_no_hole:false ())
+          else
+            None
         in
-        CompletionItem.create ~label:expr_wo_parens ~textEdit
+        CompletionItem.create ~label:expr_wo_parens ~textEdit:(`TextEdit edit)
           ~filterText:("_" ^ expr) ~kind:CompletionItemKind.Text
-          ~sortText:(sortText_of_index idx) ~command ()
+          ~sortText:(sortText_of_index idx) ?command ()
       in
       List.mapi constructed_exprs ~f:completionItem_of_constructed_expr
 end
 
-let complete doc pos =
+let complete (state : State.t)
+    ({ textDocument = { uri }; position = pos; _ } : CompletionParams.t) =
+  let doc = Document_store.get state.store uri in
   let+ items =
     let position = Position.logical pos in
     let prefix =
@@ -272,15 +279,20 @@ let complete doc pos =
             (construct_cmd_resp, compl_by_prefix_resp))
       in
       let construct_completionItems =
-        Complete_with_construct.process_dispatch_resp construct_cmd_resp
+        let supportsJumpToNextHole =
+          State.experimental_client_capabilities state
+          |> Client.Experimental_capabilities.supportsJumpToNextHole
+        in
+        Complete_with_construct.process_dispatch_resp ~supportsJumpToNextHole
+          construct_cmd_resp
       in
       let compl_by_prefix_completionItems =
         Complete_by_prefix.process_dispatch_resp doc pos compl_by_prefix_resp
       in
       construct_completionItems @ compl_by_prefix_completionItems
       |> reindex_sortText |> preselect_first
   in
-  `CompletionList (CompletionList.create ~isIncomplete:false ~items)
+  Some (`CompletionList (CompletionList.create ~isIncomplete:false ~items))
 
 let format_doc ~markdown doc =
   match markdown with

ocaml-lsp-server/src/compl.mli

@@ -13,7 +13,9 @@ module Resolve : sig
 end
 
 val complete :
-  Document.t -> Position.t -> [> `CompletionList of CompletionList.t ] Fiber.t
+     State.t
+  -> CompletionParams.t
+  -> [> `CompletionList of CompletionList.t ] option Fiber.t
 
 (** creates a server response for ["completionItem/resolve"] *)
 val resolve :

ocaml-lsp-server/src/ocaml_lsp_server.ml

@@ -837,13 +837,8 @@ let ocaml_on_request :
         definition_query rpc state uri position (fun pos ->
             Query_protocol.Locate_type pos))
       ()
-  | TextDocumentCompletion { textDocument = { uri }; position; _ } ->
-    later
-      (fun _ () ->
-        let doc = Document_store.get store uri in
-        let+ resp = Compl.complete doc position in
-        Some resp)
-      ()
+  | TextDocumentCompletion params ->
+    later (fun _ () -> Compl.complete state params) ()
   | TextDocumentPrepareRename { textDocument = { uri }; position } ->
     later
       (fun _ () ->

ocaml-lsp-server/src/range.ml

@@ -24,3 +24,21 @@ let of_loc (loc : Loc.t) : t =
   let+ end_ = Position.of_lexical_position loc.loc_end in
   { start; end_ })
   |> Option.value ~default:first_line
+
+let resize_for_edit { TextEdit.range; newText } =
+  let lines = String.split_lines newText in
+  match lines with
+  | [] -> { range with end_ = range.start }
+  | several_lines ->
+    let end_ =
+      let start = range.start in
+      let line = start.line + List.length several_lines - 1 in
+      let character =
+        let last_line_len =
+          List.last several_lines |> Option.value_exn |> String.length
+        in
+        start.character + last_line_len
+      in
+      { Position.line; character }
+    in
+    { range with end_ }

ocaml-lsp-server/src/range.mli

@@ -11,3 +11,9 @@ val compare_size : t -> t -> Ordering.t
 val first_line : t
 
 val of_loc : Loc.t -> t
+
+(** [resize_for_edit edit] returns shrunk, unchanged, or extended [edit.range]
+    depending on the size of [edit.newText], e.g., if [edit.newText] contains
+    less characters than [edit.range], the new range is shrunk to fit
+    [edit.newText] only. *)
+val resize_for_edit : TextEdit.t -> t

ocaml-lsp-server/src/state.ml

@@ -6,6 +6,7 @@ type init =
       { params : InitializeParams.t
       ; workspaces : Workspaces.t
       ; dune : Dune.t
+      ; exp_client_caps : Client.Experimental_capabilities.t
       }
 
 type t =
@@ -64,7 +65,17 @@ let dune t =
 
 let initialize t params workspaces dune =
   assert (t.init = Uninitialized);
-  { t with init = Initialized { params; workspaces; dune } }
+  { t with
+    init =
+      Initialized
+        { params
+        ; workspaces
+        ; dune
+        ; exp_client_caps =
+            Client.Experimental_capabilities.of_opt_json
+              params.capabilities.experimental
+        }
+  }
 
 let modify_workspaces t ~f =
   let init =
@@ -77,6 +88,11 @@ let modify_workspaces t ~f =
 
 let client_capabilities t = (initialize_params t).capabilities
 
+let experimental_client_capabilities t =
+  match t.init with
+  | Uninitialized -> assert false
+  | Initialized { exp_client_caps; _ } -> exp_client_caps
+
 let log_msg server ~type_ ~message =
   let state = Server.state server in
   task_if_running state.detached ~f:(fun () ->

ocaml-lsp-server/src/state.mli

@@ -6,6 +6,7 @@ type init =
       { params : InitializeParams.t
       ; workspaces : Workspaces.t
       ; dune : Dune.t
+      ; exp_client_caps : Client.Experimental_capabilities.t
       }
 
 type t =
@@ -47,9 +48,14 @@ val dune : t -> Dune.t
 
 val modify_workspaces : t -> f:(Workspaces.t -> Workspaces.t) -> t
 
-(** @return client capabilities passed from the client in [InitializeParams]
+(** @return
+      client capabilities passed from the client in [InitializeParams]; use
+      [exp_client_caps] to get {i experimental} client capabilities.
     @raise Assertion_failure if the [t.init] is [Uninitialized] *)
 val client_capabilities : t -> ClientCapabilities.t
 
+(** @return experimental client capabilities *)
+val experimental_client_capabilities : t -> Client.Experimental_capabilities.t
+
 val log_msg :
   t Server.t -> type_:MessageType.t -> message:string -> unit Fiber.t

ocaml-lsp-server/test/e2e/__tests__/textDocument-codeAction.test.ts

@@ -60,7 +60,9 @@ describe("textDocument/codeAction", () => {
   }
 
   beforeEach(async () => {
-    languageServer = await LanguageServer.startAndInitialize();
+    languageServer = await LanguageServer.startAndInitialize({
+      capabilities: { experimental: { jumpToNextHole: true } },
+    });
   });
 
   afterEach(async () => {
@@ -102,11 +104,17 @@ let f (x : t) = x
           "command": Object {
             "arguments": Array [
               Object {
-                "notify-if-no-hole": false,
-                "position": Object {
-                  "character": 16,
-                  "line": 2,
+                "inRange": Object {
+                  "end": Object {
+                    "character": 54,
+                    "line": 2,
+                  },
+                  "start": Object {
+                    "character": 16,
+                    "line": 2,
+                  },
                 },
+                "shouldNotifyIfNoHole": false,
               },
             ],
             "command": "ocaml.next-hole",