feat: waitForDiagnostics & waitForResponses requests

src/Lean/Data/Lsp.lean

@@ -8,6 +8,7 @@ import Lean.Data.Lsp.Basic
 import Lean.Data.Lsp.Capabilities
 import Lean.Data.Lsp.Communication
 import Lean.Data.Lsp.Diagnostics
+import Lean.Data.Lsp.Extra
 import Lean.Data.Lsp.Hover
 import Lean.Data.Lsp.InitShutdown
 import Lean.Data.Lsp.TextSync

src/Lean/Data/Lsp/Extra.lean

@@ -0,0 +1,64 @@
+/-
+Copyright (c) 2020 Marc Huisinga. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Authors: Marc Huisinga
+-/
+import Lean.Data.Json
+import Lean.Data.JsonRpc
+import Lean.Data.Lsp.Basic
+
+/-
+This file contains Lean-specific extensions to LSP.
+The following additional packets are supported:
+- "textDocument/waitForDiagnostics": Yields a response when all the diagnostics that were pending at the time of
+  arrival of the packet have been emitted.
+- "textDocument/waitForResponses": Yields a response when all the requests that were pending at the time of
+  arrival of the packet have been emitted.
+Both of these exist for synchronization purposes, e.g. during testing or when external tools might want to use
+our LSP server.
+-/
+
+namespace Lean.Lsp
+
+open Json
+
+structure WaitForDiagnosticsParam where
+  uri : DocumentUri
+
+instance : FromJson WaitForDiagnosticsParam :=
+  ⟨fun j => do
+    let id ← j.getObjValAs? DocumentUri "uri"
+    pure ⟨id⟩⟩
+
+instance : ToJson WaitForDiagnosticsParam :=
+  ⟨fun o => mkObj [⟨"uri", toJson o.uri⟩]⟩
+
+structure WaitForDiagnostics
+
+instance : FromJson WaitForDiagnostics :=
+  ⟨fun j => WaitForDiagnostics.mk⟩
+
+instance : ToJson WaitForDiagnostics :=
+  ⟨fun o => mkObj []⟩
+
+structure WaitForResponsesParam where
+  uri : DocumentUri
+
+instance : FromJson WaitForResponsesParam :=
+  ⟨fun j => do
+    let id ← j.getObjValAs? DocumentUri "uri"
+    pure ⟨id⟩⟩
+
+instance : ToJson WaitForResponsesParam :=
+  ⟨fun o => mkObj [⟨"uri", toJson o.uri⟩]⟩
+
+structure WaitForResponses
+
+instance : FromJson WaitForResponses :=
+  ⟨fun j => WaitForResponses.mk⟩
+
+instance : ToJson WaitForResponses :=
+  ⟨fun o => mkObj []⟩
+
+end Lean.Lsp

src/Lean/Server/AsyncList.lean

@@ -64,16 +64,29 @@ partial def unfoldAsync [Coe Error ε] (f : α → ExceptT ε IO α)
 
 /-- Waits for the entire computation to finish and returns the computed
 list. If computation was ongoing, also returns the terminating value. -/
-partial def waitAll : AsyncList ε α → List α × Option ε
+partial def getAll : AsyncList ε α → List α × Option ε
   | cons hd tl =>
-    let ⟨l, e?⟩ := tl.waitAll
+    let ⟨l, e?⟩ := tl.getAll
     ⟨hd :: l, e?⟩
   | nil => ⟨[], none⟩
   | asyncTail tl =>
     match tl.get with
-    | Except.ok tl => tl.waitAll
+    | Except.ok tl => tl.getAll
     | Except.error e => ⟨[], some e⟩
 
+/-- Waits for the entire computation to finish and returns the computed
+list. If computation was ongoing, also returns the terminating value.
+As opposed to getAll, ensures that the waiting is completed. -/
+partial def waitAll : AsyncList ε α → IO (List α × Option ε)
+  | cons hd tl => do
+    let ⟨l, e?⟩ ← tl.waitAll
+    pure ⟨hd :: l, e?⟩
+  | nil => pure ⟨[], none⟩
+  | asyncTail tl => do
+    match ←IO.wait tl with
+    | Except.ok tl => tl.waitAll
+    | Except.error e => pure ⟨[], some e⟩
+
 /-- Extends the `finishedPrefix` as far as possible. If computation was ongoing
 and has finished, also returns the terminating value. -/
 partial def updateFinishedPrefix : AsyncList ε α → IO (AsyncList ε α × Option ε)

src/Lean/Server/FileSource.lean

@@ -43,5 +43,11 @@ instance DidCloseTextDocumentParams.hasFileSource : FileSource DidCloseTextDocum
 instance HoverParams.hasFileSource : FileSource HoverParams :=
   ⟨fun h => fileSource h.toTextDocumentPositionParams⟩
 
+instance WaitForDiagnosticsParam.hasFileSource : FileSource WaitForDiagnosticsParam :=
+  ⟨fun p => p.uri⟩
+
+instance WaitForResponsesParam.hasFileSource : FileSource WaitForResponsesParam :=
+  ⟨fun p => p.uri⟩
+
 end Lsp
 end Lean

src/Lean/Server/FileWorker.lean

@@ -70,6 +70,8 @@ open Std (RBMap RBMap.empty)
 open JsonRpc
 
 section ServerM
+  -- Pending requests are tracked so that requests can be cancelled by cancelling the corresponding task,
+  -- which would be cancelled by the GC if we did not track these requests.
   abbrev PendingRequestMap := RBMap RequestID (Task (Except IO.Error Unit)) (fun a b => Decidable.decide (a < b))
 
   structure ServerContext where
@@ -203,7 +205,7 @@ section NotificationHandling
 end NotificationHandling
 
 section RequestHandling
-  /- TODO(MH): Requests that need data from a certain command should traverse e.snapshots
+  /- TODO(MH): Requests that need data from a certain command should traverse the snapshots
      by successively getting the next task, meaning that we might need to wait for elaboration.
      Sebastian said something about a future function TaskIO.bind that ensures that the
      request task will also stop waiting when the reference to the task is released by handleDidChange.
@@ -211,7 +213,23 @@ section RequestHandling
      (this way, the server doesn't get bogged down in requests for an old state of the document).
      Requests need to manually check for whether their task has been cancelled, so that they
      can reply with a RequestCancelled error. -/
-  def handleHover (p : HoverParams) (e : EditableDocument) : IO Unit := pure ⟨⟩
+  def handleHover (id : RequestID) (p : HoverParams) : ServerM Unit := pure ⟨⟩
+
+  def handleWaitForDiagnostics (id : RequestID) (p : WaitForDiagnosticsParam) : ServerM Unit := do
+    let st ← read
+    let e ← st.docRef.get
+    discard $ e.cmdSnaps.waitAll
+    st.hOut.writeLspResponse ⟨id, WaitForDiagnostics.mk⟩
+
+  /- The signature of this request is somewhat special, since it is a request that performs
+     an action related to requests. Hence we need to ensure that this handler does not
+     accidentally reference itself, by passing in the `PendingRequestMap` from before
+     the request was added. -/
+  def handleWaitForResponses (id : RequestID) (pr : PendingRequestMap) : ServerM Unit := do
+    for ⟨_, task⟩ in pr do
+      discard $ IO.wait task
+    (←read).hOut.writeLspResponse ⟨id, WaitForResponses.mk⟩
+
 end RequestHandling
 
 section MessageHandling
@@ -228,18 +246,20 @@ section MessageHandling
     | "$/cancelRequest"        => handle CancelParams handleCancelRequest
     | _                        => throwServerError $ "Got unsupported notification method: " ++ method
 
-  def queueRequest (id : RequestID) (handler : α → EditableDocument → IO Unit) (params : α)
+  def queueRequest (id : RequestID) (handler : RequestID → α → ServerM Unit) (params : α)
   : ServerM Unit := do
-    let requestTask ← asTask (handler params (←(←read).docRef.get))
+    let requestTask ← asTask (handler id params (←read))
     updatePendingRequests (fun pendingRequests => pendingRequests.insert id requestTask)
 
   def handleRequest (id : RequestID) (method : String) (params : Json)
   : ServerM Unit := do
     let handle := fun paramType [FromJson paramType]
-                      (handler : paramType → EditableDocument → IO Unit) =>
+                      (handler : RequestID → paramType → ServerM Unit) =>
       parseParams paramType params >>= queueRequest id handler
     match method with
-    | "textDocument/hover" => handle HoverParams handleHover
+    | "textDocument/waitForResponses"   => queueRequest id handleWaitForResponses (←(←read).pendingRequestsRef.get)
+    | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParam handleWaitForDiagnostics
+    | "textDocument/hover"              => handle HoverParams handleHover
     | _ => throwServerError $ "Got unsupported request: " ++ method ++
                               " params: " ++ toString params
 end MessageHandling

src/Lean/Server/Watchdog.lean

@@ -342,7 +342,9 @@ section MessageHandling
       let uri := fileSource parsedParams
       tryWriteMessage uri ⟨id, method, parsedParams⟩ FileWorker.writeRequest
     match method with
-    | "textDocument/hover" => handle HoverParams
+    | "textDocument/waitForResponses"   => handle WaitForResponsesParam
+    | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParam
+    | "textDocument/hover"              => handle HoverParams
     | _ => throwServerError $ "Got unsupported request: " ++ method ++ " params: " ++ toString params
 
   def handleNotification (method : String) (params : Json) : ServerM Unit :=
@@ -409,11 +411,11 @@ end MainLoop
 
 def mkLeanServerCapabilities : ServerCapabilities := {
   textDocumentSync? := some {
-    openClose := true
-    change := TextDocumentSyncKind.incremental
-    willSave := false
+    openClose         := true
+    change            := TextDocumentSyncKind.incremental
+    willSave          := false
     willSaveWaitUntil := false
-    save? := none
+    save?             := none
   }
   hoverProvider := true
 }