From 4e3a8468c351531235ba5cffd5566b7283b7a8ca Mon Sep 17 00:00:00 2001 From: Marc Huisinga Date: Thu, 14 Mar 2024 18:10:04 +0100 Subject: [PATCH] fix: periodically refresh semantic tokens (#3619) This PR fixes an issue where the file worker would not provide the client with semantic tokens until the file had been elaborated completely. The file worker now also tells the client to refresh its semantic tokens after running "Restart File". This PR is based on #3271. --- src/Lean/Server/FileWorker.lean | 48 ++++++++++++++----- .../Server/FileWorker/RequestHandling.lean | 35 +++++++++----- 2 files changed, 58 insertions(+), 25 deletions(-) diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean index 0cae415299..53ea69bda7 100644 --- a/src/Lean/Server/FileWorker.lean +++ b/src/Lean/Server/FileWorker.lean @@ -62,21 +62,22 @@ open JsonRpc structure WorkerContext where /-- Synchronized output channel for LSP messages. Notifications for outdated versions are discarded on read. -/ - chanOut : IO.Channel JsonRpc.Message + chanOut : IO.Channel JsonRpc.Message /-- Latest document version received by the client, used for filtering out notifications from previous versions. -/ - maxDocVersionRef : IO.Ref Int - hLog : FS.Stream - initParams : InitializeParams - processor : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot - clientHasWidgets : Bool + maxDocVersionRef : IO.Ref Int + freshRequestIdRef : IO.Ref Int + hLog : FS.Stream + initParams : InitializeParams + processor : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot + clientHasWidgets : Bool /-- Options defined on the worker cmdline (i.e. not including options from `setup-file`), used for context-free tasks such as editing delay. -/ - cmdlineOpts : Options + cmdlineOpts : Options /-! # Asynchronous snapshot elaboration -/ @@ -222,7 +223,6 @@ structure WorkerState where /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers to modify sessions. A single `Ref` ensures atomic transactions. -/ rpcSessions : RBMap UInt64 (IO.Ref RpcSession) compare - freshRequestID : Nat abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO @@ -288,6 +288,7 @@ section Initialization mainModuleName ← moduleNameOfFileName path none catch _ => pure () let maxDocVersionRef ← IO.mkRef 0 + let freshRequestIdRef ← IO.mkRef 0 let chanOut ← mkLspOutputChannel maxDocVersionRef let srcSearchPathPromise ← IO.Promise.new @@ -301,6 +302,7 @@ section Initialization processor clientHasWidgets maxDocVersionRef + freshRequestIdRef cmdlineOpts := opts } let doc : EditableDocumentCore := { @@ -316,7 +318,6 @@ section Initialization pendingRequests := RBMap.empty rpcSessions := RBMap.empty importCachingTask? := none - freshRequestID := 0 }) where /-- Creates an LSP message output channel along with a reader that sends out read messages on @@ -345,6 +346,18 @@ section Initialization return chanOut end Initialization +section ServerRequests + def sendServerRequest [ToJson α] + (ctx : WorkerContext) + (method : String) + (param : α) + : IO Unit := do + let freshRequestId ← ctx.freshRequestIdRef.modifyGet fun freshRequestId => + (freshRequestId, freshRequestId + 1) + let r : JsonRpc.Request α := ⟨freshRequestId, method, param⟩ + ctx.chanOut.send r +end ServerRequests + section Updates def updatePendingRequests (map : PendingRequestMap → PendingRequestMap) : WorkerM Unit := do modify fun st => { st with pendingRequests := map st.pendingRequests } @@ -541,8 +554,8 @@ section MessageHandling ctx.chanOut.send <| e.toLspResponseError id queueRequest id t - def handleResponse (_ : RequestID) (result : Json) : WorkerM Unit := - throwServerError s!"Unknown response kind: {result}" + def handleResponse (_ : RequestID) (_ : Json) : WorkerM Unit := + return -- The only response that we currently expect here is always empty end MessageHandling @@ -589,6 +602,13 @@ section MainLoop | _ => throwServerError "Got invalid JSON-RPC message" end MainLoop +def runRefreshTask : WorkerM (Task (Except IO.Error Unit)) := do + let ctx ← read + IO.asTask do + while ! (←IO.checkCanceled) do + sendServerRequest ctx "workspace/semanticTokens/refresh" (none : Option Nat) + IO.sleep 2000 + def initAndRunWorker (i o e : FS.Stream) (opts : Options) : IO UInt32 := do let i ← maybeTee "fwIn.txt" false i let o ← maybeTee "fwOut.txt" true o @@ -605,7 +625,11 @@ def initAndRunWorker (i o e : FS.Stream) (opts : Options) : IO UInt32 := do let _ ← IO.setStderr e try let (ctx, st) ← initializeWorker meta o e initParams.param opts - let _ ← StateRefT'.run (s := st) <| ReaderT.run (r := ctx) (mainLoop i) + let _ ← StateRefT'.run (s := st) <| ReaderT.run (r := ctx) do + let refreshTask ← runRefreshTask + let exitCode ← mainLoop i + IO.cancel refreshTask + pure exitCode return (0 : UInt32) catch err => IO.eprintln err diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean index 959ae9f5a9..467c9612b2 100644 --- a/src/Lean/Server/FileWorker/RequestHandling.lean +++ b/src/Lean/Server/FileWorker/RequestHandling.lean @@ -430,10 +430,10 @@ def noHighlightKinds : Array SyntaxNodeKind := #[ ``Lean.Parser.Command.moduleDoc] structure SemanticTokensContext where - beginPos : String.Pos - endPos : String.Pos - text : FileMap - snap : Snapshot + beginPos : String.Pos + endPos? : Option String.Pos + text : FileMap + snap : Snapshot structure SemanticTokensState where data : Array Nat @@ -447,20 +447,29 @@ def keywordSemanticTokenMap : RBMap String SemanticTokenType compare := |>.insert "stop" .leanSorryLike |>.insert "#exit" .leanSorryLike -partial def handleSemanticTokens (beginPos endPos : String.Pos) +partial def handleSemanticTokens (beginPos : String.Pos) (endPos? : Option String.Pos) : RequestM (RequestTask SemanticTokens) := do let doc ← readDoc - let text := doc.meta.text - let t := doc.cmdSnaps.waitUntil (·.endPos >= endPos) - mapTask t fun (snaps, _) => + match endPos? with + | none => + -- Only grabs the finished prefix so that we do not need to wait for elaboration to complete + -- for the full file before sending a response. This means that the response will be incomplete, + -- which we mitigate by regularly sending `workspace/semanticTokens/refresh` requests in the + -- `FileWorker` to tell the client to re-compute the semantic tokens. + let (snaps, _) ← doc.cmdSnaps.getFinishedPrefix + asTask <| run doc snaps + | some endPos => + let t := doc.cmdSnaps.waitUntil (·.endPos >= endPos) + mapTask t fun (snaps, _) => run doc snaps +where + run doc snaps : RequestM SemanticTokens := StateT.run' (s := { data := #[], lastLspPos := ⟨0, 0⟩ : SemanticTokensState }) do for s in snaps do if s.endPos <= beginPos then continue - ReaderT.run (r := SemanticTokensContext.mk beginPos endPos text s) <| + ReaderT.run (r := SemanticTokensContext.mk beginPos endPos? doc.meta.text s) <| go s.stx return { data := (← get).data } -where go (stx : Syntax) := do match stx with | `($e.$id:ident) => go e; addToken id SemanticTokenType.property @@ -506,9 +515,9 @@ where (val.length > 1 && val.front == '#' && (val.get ⟨1⟩).isAlpha) then addToken stx (keywordSemanticTokenMap.findD val .keyword) addToken stx type := do - let ⟨beginPos, endPos, text, _⟩ ← read + let ⟨beginPos, endPos?, text, _⟩ ← read if let (some pos, some tailPos) := (stx.getPos?, stx.getTailPos?) then - if beginPos <= pos && pos < endPos then + if beginPos <= pos && endPos?.all (pos < ·) then let lspPos := (← get).lastLspPos let lspPos' := text.utf8PosToLspPos pos let deltaLine := lspPos'.line - lspPos.line @@ -523,7 +532,7 @@ where def handleSemanticTokensFull (_ : SemanticTokensParams) : RequestM (RequestTask SemanticTokens) := do - handleSemanticTokens 0 ⟨1 <<< 31⟩ + handleSemanticTokens 0 none def handleSemanticTokensRange (p : SemanticTokensRangeParams) : RequestM (RequestTask SemanticTokens) := do