fix: merlin document safety (#890)

only use merlin features on merlin sources. enforce this with the type
system

Signed-off-by: Rudi Grinberg <[email protected]>

Author: rgrinberg

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

@@ -65,25 +65,28 @@ let code_action_add_rec uri diagnostics doc loc =
     ()
 
 let code_action doc (params : CodeActionParams.t) =
-  let pos_start = Position.logical params.range.start in
-  let m_diagnostic =
-    List.find params.context.diagnostics ~f:(fun d ->
-        let is_unbound () =
-          String.is_prefix d.Diagnostic.message ~prefix:"Unbound value"
-        and in_range () =
-          match Position.compare_inclusion params.range.start d.range with
-          | `Outside _ -> false
-          | `Inside -> true
-        in
-        in_range () && is_unbound ())
-  in
-  match m_diagnostic with
-  | None -> Fiber.return None
-  | Some d ->
-    let+ loc =
-      Document.with_pipeline_exn doc (fun pipeline ->
-          has_missing_rec pipeline pos_start)
+  match Document.kind doc with
+  | `Other -> Fiber.return None
+  | `Merlin merlin -> (
+    let pos_start = Position.logical params.range.start in
+    let m_diagnostic =
+      List.find params.context.diagnostics ~f:(fun d ->
+          let is_unbound () =
+            String.is_prefix d.Diagnostic.message ~prefix:"Unbound value"
+          and in_range () =
+            match Position.compare_inclusion params.range.start d.range with
+            | `Outside _ -> false
+            | `Inside -> true
+          in
+          in_range () && is_unbound ())
     in
-    Option.map loc ~f:(code_action_add_rec params.textDocument.uri [ d ] doc)
+    match m_diagnostic with
+    | None -> Fiber.return None
+    | Some d ->
+      let+ loc =
+        Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
+            has_missing_rec pipeline pos_start)
+      in
+      Option.map loc ~f:(code_action_add_rec params.textDocument.uri [ d ] doc))
 
 let t = { Code_action.kind = QuickFix; run = code_action }

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

@@ -5,8 +5,9 @@ let action_kind = "construct"
 
 let code_action doc (params : CodeActionParams.t) =
   match Document.kind doc with
-  | `Other | `Merlin Intf -> Fiber.return None
-  | `Merlin Impl ->
+  | `Other -> Fiber.return None
+  | `Merlin m when Document.Merlin.kind m = Intf -> Fiber.return None
+  | `Merlin merlin ->
     let pos = Position.logical params.range.Range.end_ in
     (* we want this predicate to quickly eliminate prefixes that don't fit to be
        a hole *)
@@ -17,7 +18,7 @@ let code_action doc (params : CodeActionParams.t) =
     if not (Typed_hole.can_be_hole prefix) then Fiber.return None
     else
       let+ structures =
-        Document.with_pipeline_exn doc (fun pipeline ->
+        Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
             let typedtree =
               let typer = Mpipeline.typer_result pipeline in
               Mtyper.get_typedtree typer

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

@@ -40,14 +40,15 @@ let code_action_of_case_analysis ~supportsJumpToNextHole doc uri (loc, newText)
 let code_action (state : State.t) doc (params : CodeActionParams.t) =
   let uri = params.textDocument.uri in
   match Document.kind doc with
-  | `Other | `Merlin Intf -> Fiber.return None
-  | `Merlin Impl -> (
+  | `Other -> Fiber.return None
+  | `Merlin m when Document.Merlin.kind m = Intf -> Fiber.return None
+  | `Merlin merlin -> (
     let command =
       let start = Position.logical params.range.start in
       let finish = Position.logical params.range.end_ in
       Query_protocol.Case_analysis (start, finish)
     in
-    let* res = Document.dispatch doc command in
+    let* res = Document.Merlin.dispatch merlin command in
     match res with
     | Ok (loc, newText) ->
       let+ newText =

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

@@ -26,8 +26,9 @@ let code_action_of_intf doc intf range =
 
 let code_action (state : State.t) doc (params : CodeActionParams.t) =
   match Document.kind doc with
-  | `Other | `Merlin Impl -> Fiber.return None
-  | `Merlin Intf -> (
+  | `Other -> Fiber.return None
+  | `Merlin m when Document.Merlin.kind m = Impl -> Fiber.return None
+  | `Merlin _ -> (
     let* intf = Inference.infer_intf state doc in
     match intf with
     | None -> Fiber.return None

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

@@ -403,53 +403,58 @@ let inline_edits pipeline task =
 
 let code_action doc (params : CodeActionParams.t) =
   let open Option.O in
-  Document.with_pipeline_exn doc (fun pipeline ->
-      let* typedtree =
-        match Mtyper.get_typedtree (Mpipeline.typer_result pipeline) with
-        | `Interface _ -> None
-        | `Implementation x -> Some x
-      in
-      let* task = find_inline_task typedtree params.range.start in
-      inline_edits pipeline task)
-  |> Fiber.map ~f:(fun m_edits ->
-         let* edits, m_error = m_edits in
-         match (edits, m_error) with
-         | [], None -> None
-         | [], Some error ->
-           let action =
-             CodeAction.create
-               ~title:action_title
-               ~kind:CodeActionKind.RefactorInline
-               ~isPreferred:false
-               ~disabled:
-                 (CodeAction.create_disabled ~reason:(string_of_error error))
-               ()
-           in
-           Some action
-         | _ :: _, (Some _ | None) ->
-           let edit =
-             let version = Document.version doc in
-             let textDocument =
-               OptionalVersionedTextDocumentIdentifier.create
-                 ~uri:params.textDocument.uri
-                 ~version
+  match Document.kind doc with
+  | `Other -> Fiber.return None
+  | `Merlin merlin ->
+    Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
+        let* typedtree =
+          match Mtyper.get_typedtree (Mpipeline.typer_result pipeline) with
+          | `Interface _ -> None
+          | `Implementation x -> Some x
+        in
+        let* task = find_inline_task typedtree params.range.start in
+        inline_edits pipeline task)
+    |> Fiber.map ~f:(fun m_edits ->
+           let* edits, m_error = m_edits in
+           match (edits, m_error) with
+           | [], None -> None
+           | [], Some error ->
+             let action =
+               CodeAction.create
+                 ~title:action_title
+                 ~kind:CodeActionKind.RefactorInline
+                 ~isPreferred:false
+                 ~disabled:
+                   (CodeAction.create_disabled ~reason:(string_of_error error))
                  ()
              in
+             Some action
+           | _ :: _, (Some _ | None) ->
              let edit =
-               TextDocumentEdit.create
-                 ~textDocument
-                 ~edits:(List.map edits ~f:(fun e -> `TextEdit e))
+               let version = Document.version doc in
+               let textDocument =
+                 OptionalVersionedTextDocumentIdentifier.create
+                   ~uri:params.textDocument.uri
+                   ~version
+                   ()
+               in
+               let edit =
+                 TextDocumentEdit.create
+                   ~textDocument
+                   ~edits:(List.map edits ~f:(fun e -> `TextEdit e))
+               in
+               WorkspaceEdit.create
+                 ~documentChanges:[ `TextDocumentEdit edit ]
+                 ()
+             in
+             let action =
+               CodeAction.create
+                 ~title:action_title
+                 ~kind:CodeActionKind.RefactorInline
+                 ~edit
+                 ~isPreferred:false
+                 ()
              in
-             WorkspaceEdit.create ~documentChanges:[ `TextDocumentEdit edit ] ()
-           in
-           let action =
-             CodeAction.create
-               ~title:action_title
-               ~kind:CodeActionKind.RefactorInline
-               ~edit
-               ~isPreferred:false
-               ()
-           in
-           Some action)
+             Some action)
 
 let t = { Code_action.kind = RefactorInline; run = code_action }

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

@@ -71,7 +71,7 @@ let rec mark_value_unused_edit name contexts =
 
 let code_action_mark_value_unused doc (diagnostic : Diagnostic.t) =
   let open Option.O in
-  Document.with_pipeline_exn doc (fun pipeline ->
+  Document.Merlin.with_pipeline_exn (Document.merlin_exn doc) (fun pipeline ->
       let var_name = slice doc diagnostic.range in
       let pos = diagnostic.range.start in
       let+ text_edit =
@@ -125,10 +125,9 @@ let code_action_remove_range doc (diagnostic : Diagnostic.t) range =
 
 (* Create a code action that removes the value mentioned in [diagnostic]. *)
 let code_action_remove_value doc pos (diagnostic : Diagnostic.t) =
-  Document.with_pipeline_exn doc (fun pipeline ->
+  Document.Merlin.with_pipeline_exn (Document.merlin_exn doc) (fun pipeline ->
       let var_name = slice doc diagnostic.range in
-      enclosing_pos pipeline pos
-      |> List.map ~f:(fun (_, x) -> x)
+      enclosing_pos pipeline pos |> List.map ~f:snd
       |> enclosing_value_binding_range var_name
       |> Option.map ~f:(fun range ->
              code_action_remove_range doc diagnostic range))

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

@@ -3,30 +3,33 @@ open Fiber.O
 
 let code_action (mode : [ `Qualify | `Unqualify ]) (action_kind : string) doc
     (params : CodeActionParams.t) =
-  let+ res =
-    let command =
-      let pos_start = Position.logical params.range.start in
-      Query_protocol.Refactor_open (mode, pos_start)
+  match Document.kind doc with
+  | `Other -> Fiber.return None
+  | `Merlin doc -> (
+    let+ res =
+      let command =
+        let pos_start = Position.logical params.range.start in
+        Query_protocol.Refactor_open (mode, pos_start)
+      in
+      Document.Merlin.dispatch_exn doc command
     in
-    Document.dispatch_exn doc command
-  in
-  match res with
-  | [] -> None
-  | changes ->
-    let code_action =
-      let edit : WorkspaceEdit.t =
-        let edits =
-          List.map changes ~f:(fun (newText, loc) ->
-              { TextEdit.newText; range = Range.of_loc loc })
+    match res with
+    | [] -> None
+    | changes ->
+      let code_action =
+        let edit : WorkspaceEdit.t =
+          let edits =
+            List.map changes ~f:(fun (newText, loc) ->
+                { TextEdit.newText; range = Range.of_loc loc })
+          in
+          let uri = params.textDocument.uri in
+          WorkspaceEdit.create ~changes:[ (uri, edits) ] ()
         in
-        let uri = params.textDocument.uri in
-        WorkspaceEdit.create ~changes:[ (uri, edits) ] ()
+        let kind = CodeActionKind.Other action_kind in
+        let title = String.capitalize_ascii action_kind in
+        CodeAction.create ~title ~kind ~edit ~isPreferred:false ()
       in
-      let kind = CodeActionKind.Other action_kind in
-      let title = String.capitalize_ascii action_kind in
-      CodeAction.create ~title ~kind ~edit ~isPreferred:false ()
-    in
-    Some code_action
+      Some code_action)
 
 let unqualify =
   let action_kind = "remove module name from identifiers" in

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

@@ -44,25 +44,31 @@ let code_action_of_type_enclosing uri doc (loc, typ) =
     ()
 
 let code_action doc (params : CodeActionParams.t) =
-  let pos_start = Position.logical params.range.start in
-  let+ res =
-    Document.with_pipeline_exn doc (fun pipeline ->
-        let context = check_typeable_context pipeline pos_start in
-        match context with
-        | `Invalid -> None
-        | `Valid ->
-          let command = Query_protocol.Type_enclosing (None, pos_start, None) in
-          let config = Mpipeline.final_config pipeline in
-          let config =
-            { config with query = { config.query with verbosity = 0 } }
-          in
-          let pipeline = Mpipeline.make config (Document.source doc) in
-          Some (Query_commands.dispatch pipeline command))
-  in
-  match res with
-  | None | Some [] | Some ((_, `Index _, _) :: _) -> None
-  | Some ((location, `String value, _) :: _) ->
-    code_action_of_type_enclosing params.textDocument.uri doc (location, value)
+  match Document.kind doc with
+  | `Other -> Fiber.return None
+  | `Merlin merlin -> (
+    let pos_start = Position.logical params.range.start in
+    let+ res =
+      Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
+          let context = check_typeable_context pipeline pos_start in
+          match context with
+          | `Invalid -> None
+          | `Valid ->
+            let command =
+              Query_protocol.Type_enclosing (None, pos_start, None)
+            in
+            let config = Mpipeline.final_config pipeline in
+            let config =
+              { config with query = { config.query with verbosity = 0 } }
+            in
+            let pipeline = Mpipeline.make config (Document.source doc) in
+            Some (Query_commands.dispatch pipeline command))
+    in
+    match res with
+    | None | Some [] | Some ((_, `Index _, _) :: _) -> None
+    | Some ((location, `String value, _) :: _) ->
+      code_action_of_type_enclosing params.textDocument.uri doc (location, value)
+    )
 
 let t =
   { Code_action.kind = CodeActionKind.Other action_kind; run = code_action }