From 829f1b286645c27e919a9b2e7f3de3fc12c60475 Mon Sep 17 00:00:00 2001
From: Rudi Grinberg <me@rgrinberg.com>
Date: Mon, 24 Apr 2023 11:29:17 -0600
Subject: [PATCH] fix: allow opening the same document twice

this is done to accommodate neovim's lsp client

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>
---
 CHANGES.md                                    |  3 ++
 ocaml-lsp-server/src/document_store.ml        |  8 +++-
 ocaml-lsp-server/src/ocaml_lsp_server.ml      |  1 -
 .../test/e2e-new/document_flow.ml             | 47 +++++++------------
 4 files changed, 25 insertions(+), 34 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 28c079c4a..b9e7a4f1b 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -2,6 +2,9 @@
 
 ## Fixes
 
+- Allow opening documents that were already open. This is a workaround for
+  neovim's lsp client (#1067)
+
 - Disable type annotation for functions (#1054)
 
 - Respect codeActionLiteralSupport capability (#1046)
diff --git a/ocaml-lsp-server/src/document_store.ml b/ocaml-lsp-server/src/document_store.ml
index 2c188b5b0..fdc7ede06 100644
--- a/ocaml-lsp-server/src/document_store.ml
+++ b/ocaml-lsp-server/src/document_store.ml
@@ -106,9 +106,13 @@ let open_document t doc =
          { document = Some doc; promotions = 0; semantic_tokens_cache = None });
     Fiber.return ()
   | Some d ->
-    assert (!d.document = None);
+    (* if there's no document, then we just opened it to track promotions.
+
+       if there's a document already, we're doing a double open and there's no
+       need to unregister. *)
+    let unregister = !d.document = None in
     d := { !d with document = Some doc };
-    unregister_request t [ key ]
+    if unregister then unregister_request t [ key ] else Fiber.return ()
 
 let get_opt t uri = Table.find t.db uri |> Option.bind ~f:(fun d -> !d.document)
 
diff --git a/ocaml-lsp-server/src/ocaml_lsp_server.ml b/ocaml-lsp-server/src/ocaml_lsp_server.ml
index 9e65259ef..479d9ca63 100644
--- a/ocaml-lsp-server/src/ocaml_lsp_server.ml
+++ b/ocaml-lsp-server/src/ocaml_lsp_server.ml
@@ -683,7 +683,6 @@ let on_notification server (notification : Client_notification.t) :
         state.merlin
         params
     in
-    assert (Document_store.get_opt store params.textDocument.uri = None);
     let* () = Document_store.open_document store doc in
     let+ () = set_diagnostics state.detached (State.diagnostics state) doc in
     state
diff --git a/ocaml-lsp-server/test/e2e-new/document_flow.ml b/ocaml-lsp-server/test/e2e-new/document_flow.ml
index 298574301..add605915 100644
--- a/ocaml-lsp-server/test/e2e-new/document_flow.ml
+++ b/ocaml-lsp-server/test/e2e-new/document_flow.ml
@@ -1,13 +1,24 @@
 open Test.Import
 
 let%expect_test "it should allow double opening the same document" =
-  let diagnostics = Fiber.Ivar.create () in
+  let diagnostics = Fiber.Mvar.create () in
+  let drain_diagnostics () = Fiber.Mvar.read diagnostics in
   let handler =
+    let on_request (type resp state) (client : state Client.t)
+        (req : resp Lsp.Server_request.t) :
+        (resp Lsp_fiber.Rpc.Reply.t * state) Fiber.t =
+      match req with
+      | Lsp.Server_request.ClientUnregisterCapability _ ->
+        let state = Client.state client in
+        Fiber.return (Lsp_fiber.Rpc.Reply.now (), state)
+      | _ -> assert false
+    in
     Client.Handler.make
       ~on_notification:
         (fun _ -> function
-          | PublishDiagnostics _ -> Fiber.Ivar.fill diagnostics ()
+          | PublishDiagnostics _ -> Fiber.Mvar.write diagnostics ()
           | _ -> Fiber.return ())
+      ~on_request:{ Client.Handler.on_request }
       ()
   in
   ( Test.run ~handler @@ fun client ->
@@ -35,36 +46,10 @@ let%expect_test "it should allow double opening the same document" =
           (TextDocumentDidOpen (DidOpenTextDocumentParams.create ~textDocument))
       in
       let* () = open_ "text 1" in
+      let* () = drain_diagnostics () in
       let+ () = open_ "text 2" in
       ()
     in
     Fiber.fork_and_join_unit run_client (fun () ->
-        run >>> Fiber.Ivar.read diagnostics >>> Client.stop client) );
-  [%expect
-    {|
-    (* CR expect_test_collector: This test expectation appears to contain a backtrace.
-       This is strongly discouraged as backtraces are fragile.
-       Please change this test to not include a backtrace. *)
-
-    Uncaught error when handling notification:
-    {
-      "params": {
-        "textDocument": {
-          "languageId": "ocaml",
-          "text": "text 2",
-          "uri": "file:///foo.ml",
-          "version": 0
-        }
-      },
-      "method": "textDocument/didOpen",
-      "jsonrpc": "2.0"
-    }
-    Error:
-    [ { exn =
-          "File \"ocaml-lsp-server/src/ocaml_lsp_server.ml\", line 686, characters 4-10: Assertion failed"
-      ; backtrace =
-          "Raised at Ocaml_lsp_server.on_notification in file \"ocaml-lsp-server/src/ocaml_lsp_server.ml\", line 686, characters 4-72\n\
-           Called from Fiber__Scheduler.exec in file \"fiber/src/scheduler.ml\", line 73, characters 8-11\n\
-           "
-      }
-    ] |}]
+        run >>> drain_diagnostics () >>> Client.stop client) );
+  [%expect {| |}]