| /- | |
| Copyright (c) 2020 Marc Huisinga. All rights reserved. | |
| Released under Apache 2.0 license as described in the file LICENSE. | |
| Authors: Marc Huisinga, Wojciech Nawrocki | |
| -/ | |
| import Init.System.IO | |
| import Std.Data.RBMap | |
| import Lean.Environment | |
| import Lean.PrettyPrinter | |
| import Lean.DeclarationRange | |
| import Lean.Data.Lsp | |
| import Lean.Data.Json.FromToJson | |
| import Lean.Server.Snapshots | |
| import Lean.Server.Utils | |
| import Lean.Server.AsyncList | |
| import Lean.Server.InfoUtils | |
| import Lean.Server.Completion | |
| /-! | |
| For general server architecture, see `README.md`. For details of IPC communication, see `Watchdog.lean`. | |
| This module implements per-file worker processes. | |
| File processing and requests+notifications against a file should be concurrent for two reasons: | |
| - By the LSP standard, requests should be cancellable. | |
| - Since Lean allows arbitrary user code to be executed during elaboration via the tactic framework, | |
| elaboration can be extremely slow and even not halt in some cases. Users should be able to | |
| work with the file while this is happening, e.g. make new changes to the file or send requests. | |
| To achieve these goals, elaboration is executed in a chain of tasks, where each task corresponds to | |
| the elaboration of one command. When the elaboration of one command is done, the next task is spawned. | |
| On didChange notifications, we search for the task in which the change occured. If we stumble across | |
| a task that has not yet finished before finding the task we're looking for, we terminate it | |
| and start the elaboration there, otherwise we start the elaboration at the task where the change occured. | |
| Requests iterate over tasks until they find the command that they need to answer the request. | |
| In order to not block the main thread, this is done in a request task. | |
| If a task that the request task waits for is terminated, a change occured somewhere before the | |
| command that the request is looking for and the request sends a "content changed" error. | |
| -/ | |
| namespace Lean.Server.FileWorker | |
| open Lsp | |
| open IO | |
| open Snapshots | |
| open Lean.Parser.Command | |
| open Std (RBMap RBMap.empty) | |
| open JsonRpc | |
| section Utils | |
| private def logSnapContent (s : Snapshot) (text : FileMap) : IO Unit := | |
| IO.eprintln s!"[{s.beginPos}, {s.endPos}]: `{text.source.extract s.beginPos (s.endPos-1)}`" | |
| inductive ElabTaskError where | |
| | aborted | |
| | eof | |
| | ioError (e : IO.Error) | |
| instance : Coe IO.Error ElabTaskError := ⟨ElabTaskError.ioError⟩ | |
| structure CancelToken where | |
| ref : IO.Ref Bool | |
| deriving Inhabited | |
| namespace CancelToken | |
| def new : IO CancelToken := | |
| CancelToken.mk <$> IO.mkRef false | |
| def check [MonadExceptOf ElabTaskError m] [MonadLiftT (ST RealWorld) m] [Monad m] (tk : CancelToken) : m Unit := do | |
| let c ← tk.ref.get | |
| if c then | |
| throw ElabTaskError.aborted | |
| def set (tk : CancelToken) : IO Unit := | |
| tk.ref.set true | |
| end CancelToken | |
| /-- A document editable in the sense that we track the environment | |
| and parser state after each command so that edits can be applied | |
| without recompiling code appearing earlier in the file. -/ | |
| structure EditableDocument where | |
| meta : DocumentMeta | |
| /- The first snapshot is that after the header. -/ | |
| headerSnap : Snapshot | |
| /- Subsequent snapshots occur after each command. -/ | |
| cmdSnaps : AsyncList ElabTaskError Snapshot | |
| cancelTk : CancelToken | |
| deriving Inhabited | |
| end Utils | |
| /- Asynchronous snapshot elaboration. -/ | |
| section Elab | |
| def publishDiagnostics (m : DocumentMeta) (diagnostics : Array Lsp.Diagnostic) (hOut : FS.Stream) : IO Unit := | |
| hOut.writeLspNotification { | |
| method := "textDocument/publishDiagnostics" | |
| param := { | |
| uri := m.uri | |
| version? := m.version | |
| diagnostics := diagnostics | |
| : PublishDiagnosticsParams | |
| } | |
| } | |
| def publishMessages (m : DocumentMeta) (msgLog : MessageLog) (hOut : FS.Stream) : IO Unit := do | |
| let diagnostics ← msgLog.msgs.mapM (msgToDiagnostic m.text) | |
| publishDiagnostics m diagnostics.toArray hOut | |
| /-- Elaborates the next command after `parentSnap` and emits diagnostics into `hOut`. -/ | |
| private def nextCmdSnap (m : DocumentMeta) (parentSnap : Snapshot) (cancelTk : CancelToken) (hOut : FS.Stream) | |
| : ExceptT ElabTaskError IO Snapshot := do | |
| cancelTk.check | |
| let maybeSnap ← compileNextCmd m.text.source parentSnap | |
| -- TODO(MH): check for interrupt with increased precision | |
| cancelTk.check | |
| match maybeSnap with | |
| | Sum.inl snap => | |
| /- NOTE(MH): This relies on the client discarding old diagnostics upon receiving new ones | |
| while prefering newer versions over old ones. The former is necessary because we do | |
| not explicitly clear older diagnostics, while the latter is necessary because we do | |
| not guarantee that diagnostics are emitted in order. Specifically, it may happen that | |
| we interrupted this elaboration task right at this point and a newer elaboration task | |
| emits diagnostics, after which we emit old diagnostics because we did not yet detect | |
| the interrupt. Explicitly clearing diagnostics is difficult for a similar reason, | |
| because we cannot guarantee that no further diagnostics are emitted after clearing | |
| them. -/ | |
| publishMessages m (snap.msgLog.add { | |
| fileName := "<ignored>" | |
| pos := m.text.toPosition snap.endPos | |
| severity := MessageSeverity.information | |
| data := "processing..." | |
| }) hOut | |
| snap | |
| | Sum.inr msgLog => | |
| publishMessages m msgLog hOut | |
| throw ElabTaskError.eof | |
| /-- Elaborates all commands after `initSnap`, emitting the diagnostics into `hOut`. -/ | |
| def unfoldCmdSnaps (m : DocumentMeta) (initSnap : Snapshot) (cancelTk : CancelToken) (hOut : FS.Stream) | |
| (initial : Bool) : | |
| IO (AsyncList ElabTaskError Snapshot) := do | |
| if initial && initSnap.msgLog.hasErrors then | |
| -- treat header processing errors as fatal so users aren't swamped with followup errors | |
| AsyncList.nil | |
| else | |
| AsyncList.unfoldAsync (nextCmdSnap m . cancelTk hOut) initSnap | |
| end Elab | |
| -- Pending requests are tracked so they can be cancelled | |
| abbrev PendingRequestMap := RBMap RequestID (Task (Except IO.Error Unit)) compare | |
| structure ServerContext where | |
| hIn : FS.Stream | |
| hOut : FS.Stream | |
| hLog : FS.Stream | |
| srcSearchPath : SearchPath | |
| docRef : IO.Ref EditableDocument | |
| pendingRequestsRef : IO.Ref PendingRequestMap | |
| abbrev ServerM := ReaderT ServerContext IO | |
| /- Worker initialization sequence. -/ | |
| section Initialization | |
| /-- Use `leanpkg print-paths` to compile dependencies on the fly and add them to `LEAN_PATH`. | |
| Compilation progress is reported to `hOut` via LSP notifications. Return the search path for | |
| source files. -/ | |
| partial def leanpkgSetupSearchPath (leanpkgPath : String) (m : DocumentMeta) (imports : Array Import) (hOut : FS.Stream) : IO SearchPath := do | |
| let leanpkgProc ← Process.spawn { | |
| stdin := Process.Stdio.null | |
| stdout := Process.Stdio.piped | |
| stderr := Process.Stdio.piped | |
| cmd := leanpkgPath | |
| args := #["print-paths"] ++ imports.map (toString ·.module) | |
| } | |
| -- progress notification: report latest stderr line | |
| let rec processStderr (acc : String) : IO String := do | |
| let line ← leanpkgProc.stderr.getLine | |
| if line == "" then | |
| return acc | |
| else | |
| publishDiagnostics m #[{ range := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩, severity? := DiagnosticSeverity.information, message := line }] hOut | |
| processStderr (acc ++ line) | |
| let stderr ← IO.asTask (processStderr "") Task.Priority.dedicated | |
| let stdout := String.trim (← leanpkgProc.stdout.readToEnd) | |
| let stderr ← IO.ofExcept stderr.get | |
| if (← leanpkgProc.wait) == 0 then | |
| match stdout.split (· == '\n') with | |
| | [""] => pure [] -- e.g. no leanpkg.toml | |
| | [leanPath, leanSrcPath] => let sp ← getBuiltinSearchPath | |
| let sp ← addSearchPathFromEnv sp | |
| let sp ← parseSearchPath leanPath sp | |
| searchPathRef.set sp | |
| let srcPath := parseSearchPath leanSrcPath | |
| srcPath.mapM realPathNormalized | |
| | _ => throw <| IO.userError s!"unexpected output from `leanpkg print-paths`:\n{stdout}\nstderr:\n{stderr}" | |
| else | |
| throw <| IO.userError s!"`leanpkg print-paths` failed:\n{stdout}\nstderr:\n{stderr}" | |
| def compileHeader (m : DocumentMeta) (hOut : FS.Stream) : IO (Snapshot × SearchPath) := do | |
| let opts := {} -- TODO | |
| let inputCtx := Parser.mkInputContext m.text.source "<input>" | |
| let (headerStx, headerParserState, msgLog) ← Parser.parseHeader inputCtx | |
| let leanpkgPath ← match ← IO.getEnv "LEAN_SYSROOT" with | |
| | some path => s!"{path}/bin/leanpkg{System.FilePath.exeSuffix}" | |
| | _ => s!"{← appDir}/leanpkg{System.FilePath.exeSuffix}" | |
| let mut srcSearchPath := [s!"{← appDir}/../lib/lean/src"] | |
| if let some p := (← IO.getEnv "LEAN_SRC_PATH") then | |
| srcSearchPath := srcSearchPath ++ parseSearchPath p | |
| let (headerEnv, msgLog) ← try | |
| -- NOTE: leanpkg does not exist in stage 0 (yet?) | |
| if (← fileExists leanpkgPath) then | |
| let pkgSearchPath ← leanpkgSetupSearchPath leanpkgPath m (Lean.Elab.headerToImports headerStx).toArray hOut | |
| srcSearchPath := srcSearchPath ++ pkgSearchPath | |
| Elab.processHeader headerStx opts msgLog inputCtx | |
| catch e => -- should be from `leanpkg print-paths` | |
| let msgs := MessageLog.empty.add { fileName := "<ignored>", pos := ⟨0, 0⟩, data := e.toString } | |
| publishMessages m msgs hOut | |
| pure (← mkEmptyEnvironment, msgs) | |
| let cmdState := Elab.Command.mkState headerEnv msgLog opts | |
| let cmdState := { cmdState with infoState.enabled := true, scopes := [{ header := "", opts := opts }] } | |
| let headerSnap := { | |
| beginPos := 0 | |
| stx := headerStx | |
| mpState := headerParserState | |
| cmdState := cmdState | |
| } | |
| return (headerSnap, srcSearchPath) | |
| def initializeWorker (meta : DocumentMeta) (i o e : FS.Stream) | |
| : IO ServerContext := do | |
| /- NOTE(WN): `toFileMap` marks line beginnings as immediately following | |
| "\n", which should be enough to handle both LF and CRLF correctly. | |
| This is because LSP always refers to characters by (line, column), | |
| so if we get the line number correct it shouldn't matter that there | |
| is a CR there. -/ | |
| let (headerSnap, srcSearchPath) ← compileHeader meta o | |
| let cancelTk ← CancelToken.new | |
| let cmdSnaps ← unfoldCmdSnaps meta headerSnap cancelTk o (initial := true) | |
| let doc : EditableDocument := ⟨meta, headerSnap, cmdSnaps, cancelTk⟩ | |
| return { | |
| hIn := i | |
| hOut := o | |
| hLog := e | |
| srcSearchPath := srcSearchPath | |
| docRef := ←IO.mkRef doc | |
| pendingRequestsRef := ←IO.mkRef RBMap.empty | |
| } | |
| end Initialization | |
| section Updates | |
| def updatePendingRequests (map : PendingRequestMap → PendingRequestMap) : ServerM Unit := do | |
| (←read).pendingRequestsRef.modify map | |
| /-- Given the new document and `changePos`, the UTF-8 offset of a change into the pre-change source, | |
| updates editable doc state. -/ | |
| def updateDocument (newMeta : DocumentMeta) (changePos : String.Pos) : ServerM Unit := do | |
| -- The watchdog only restarts the file worker when the syntax tree of the header changes. | |
| -- If e.g. a newline is deleted, it will not restart this file worker, but we still | |
| -- need to reparse the header so that the offsets are correct. | |
| let st ← read | |
| let oldDoc ← st.docRef.get | |
| let newHeaderSnap ← reparseHeader newMeta.text.source oldDoc.headerSnap | |
| if newHeaderSnap.stx != oldDoc.headerSnap.stx then | |
| throwServerError "Internal server error: header changed but worker wasn't restarted." | |
| let ⟨cmdSnaps, e?⟩ ← oldDoc.cmdSnaps.updateFinishedPrefix | |
| match e? with | |
| -- This case should not be possible. only the main task aborts tasks and ensures that aborted tasks | |
| -- do not show up in `snapshots` of an EditableDocument. | |
| | some ElabTaskError.aborted => | |
| throwServerError "Internal server error: elab task was aborted while still in use." | |
| | some (ElabTaskError.ioError ioError) => throw ioError | |
| | _ => -- No error or EOF | |
| oldDoc.cancelTk.set | |
| -- NOTE(WN): we invalidate eagerly as `endPos` consumes input greedily. To re-elaborate only | |
| -- when really necessary, we could do a whitespace-aware `Syntax` comparison instead. | |
| let mut validSnaps := cmdSnaps.finishedPrefix.takeWhile (fun s => s.endPos < changePos) | |
| if validSnaps.length = 0 then | |
| let cancelTk ← CancelToken.new | |
| let newCmdSnaps ← unfoldCmdSnaps newMeta newHeaderSnap cancelTk st.hOut (initial := true) | |
| st.docRef.set ⟨newMeta, newHeaderSnap, newCmdSnaps, cancelTk⟩ | |
| else | |
| /- When at least one valid non-header snap exists, it may happen that a change does not fall | |
| within the syntactic range of that last snap but still modifies it by appending tokens. | |
| We check for this here. We do not currently handle crazy grammars in which an appended | |
| token can merge two or more previous commands into one. To do so would require reparsing | |
| the entire file. -/ | |
| let mut lastSnap := validSnaps.getLast! | |
| let preLastSnap := | |
| if validSnaps.length ≥ 2 | |
| then validSnaps.get! (validSnaps.length - 2) | |
| else newHeaderSnap | |
| let newLastStx ← parseNextCmd newMeta.text.source preLastSnap | |
| if newLastStx != lastSnap.stx then | |
| validSnaps ← validSnaps.dropLast | |
| lastSnap ← preLastSnap | |
| let cancelTk ← CancelToken.new | |
| let newSnaps ← unfoldCmdSnaps newMeta lastSnap cancelTk st.hOut (initial := false) | |
| let newCmdSnaps := AsyncList.ofList validSnaps ++ newSnaps | |
| st.docRef.set ⟨newMeta, newHeaderSnap, newCmdSnaps, cancelTk⟩ | |
| end Updates | |
| /- Notifications are handled in the main thread. They may change global worker state | |
| such as the current file contents. -/ | |
| section NotificationHandling | |
| def handleDidChange (p : DidChangeTextDocumentParams) : ServerM Unit := do | |
| let docId := p.textDocument | |
| let changes := p.contentChanges | |
| let oldDoc ← (←read).docRef.get | |
| let some newVersion ← pure docId.version? | |
| | throwServerError "Expected version number" | |
| if newVersion ≤ oldDoc.meta.version then | |
| -- TODO(WN): This happens on restart sometimes. | |
| IO.eprintln s!"Got outdated version number: {newVersion} ≤ {oldDoc.meta.version}" | |
| else if ¬ changes.isEmpty then | |
| let (newDocText, minStartOff) := foldDocumentChanges changes oldDoc.meta.text | |
| updateDocument ⟨docId.uri, newVersion, newDocText⟩ minStartOff | |
| def handleCancelRequest (p : CancelParams) : ServerM Unit := do | |
| updatePendingRequests (fun pendingRequests => pendingRequests.erase p.id) | |
| end NotificationHandling | |
| /- Request handlers are given by `Task`s executed asynchronously. They may be cancelled at any time, | |
| so they should check the cancellation token when possible to handle this cooperatively. Any exceptions | |
| thrown in a handler will be reported to the client as LSP error responses. -/ | |
| section RequestHandling | |
| structure RequestError where | |
| code : ErrorCode | |
| message : String | |
| namespace RequestError | |
| def fileChanged : RequestError := | |
| { code := ErrorCode.contentModified | |
| message := "File changed." } | |
| end RequestError | |
| -- TODO(WN): the type is too complicated | |
| abbrev RequestM α := ServerM $ Task $ Except IO.Error $ Except RequestError α | |
| def mapTask (t : Task α) (f : α → ExceptT RequestError ServerM β) : RequestM β := fun st => | |
| (IO.mapTask · t) fun a => f a st | |
| /-- Create a task which waits for a snapshot matching `p`, handles various errors, | |
| and if a matching snapshot was found executes `x` with it. If not found, the task | |
| executes `notFoundX`. -/ | |
| def withWaitFindSnap (doc : EditableDocument) (p : Snapshot → Bool) | |
| (notFoundX : ExceptT RequestError ServerM β) | |
| (x : Snapshot → ExceptT RequestError ServerM β) | |
| : ServerM (Task (Except IO.Error (Except RequestError β))) := do | |
| let findTask ← doc.cmdSnaps.waitFind? p | |
| mapTask findTask fun | |
| | Except.error ElabTaskError.aborted => | |
| throwThe RequestError RequestError.fileChanged | |
| | Except.error (ElabTaskError.ioError e) => | |
| throwThe IO.Error e | |
| | Except.error ElabTaskError.eof => notFoundX | |
| | Except.ok none => notFoundX | |
| | Except.ok (some snap) => x snap | |
| /- 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. | |
| When that happens, the request should send a "content changed" error to the user | |
| (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. -/ | |
| partial def handleCompletion (p : CompletionParams) : | |
| ServerM (Task (Except IO.Error (Except RequestError CompletionList))) := do | |
| let st ← read | |
| let doc ← st.docRef.get | |
| let text := doc.meta.text | |
| let pos := text.lspPosToUtf8Pos p.position | |
| -- dbg_trace ">> handleCompletion invoked {pos}" | |
| withWaitFindSnap doc (fun s => s.endPos > pos) | |
| (notFoundX := pure { items := #[], isIncomplete := true }) fun snap => do | |
| for infoTree in snap.cmdState.infoState.trees do | |
| -- for (ctx, info) in infoTree.getCompletionInfos do | |
| -- dbg_trace "{← info.format ctx}" | |
| if let some r ← Completion.find? doc.meta.text pos infoTree then | |
| return r | |
| return { items := #[ ], isIncomplete := true } | |
| open Elab in | |
| partial def handleHover (p : HoverParams) | |
| : ServerM (Task (Except IO.Error (Except RequestError (Option Hover)))) := do | |
| let st ← read | |
| let doc ← st.docRef.get | |
| let text := doc.meta.text | |
| let mkHover (s : String) (f : String.Pos) (t : String.Pos) : Hover := | |
| { contents := { kind := MarkupKind.markdown | |
| value := s } | |
| range? := some { start := text.utf8PosToLspPos f | |
| «end» := text.utf8PosToLspPos t } } | |
| let hoverPos := text.lspPosToUtf8Pos p.position | |
| withWaitFindSnap doc (fun s => s.endPos > hoverPos) | |
| (notFoundX := pure none) fun snap => do | |
| for t in snap.cmdState.infoState.trees do | |
| if let some (ci, i) := t.hoverableInfoAt? hoverPos then | |
| if let some hoverFmt ← i.fmtHover? ci then | |
| return some <| mkHover (toString hoverFmt) i.pos?.get! i.tailPos?.get! | |
| return none | |
| open Elab in | |
| partial def handleDefinition (goToType? : Bool) (p : TextDocumentPositionParams) | |
| : ServerM (Task (Except IO.Error (Except RequestError (Array LocationLink)))) := do | |
| let st ← read | |
| let doc ← st.docRef.get | |
| let text := doc.meta.text | |
| let hoverPos := text.lspPosToUtf8Pos p.position | |
| withWaitFindSnap doc (fun s => s.endPos > hoverPos) | |
| (notFoundX := pure #[]) fun snap => do | |
| for t in snap.cmdState.infoState.trees do | |
| if let some (ci, Info.ofTermInfo i) := t.hoverableInfoAt? hoverPos then | |
| let mut expr := i.expr | |
| if goToType? then | |
| expr ← ci.runMetaM i.lctx do | |
| Meta.instantiateMVars (← Meta.inferType expr) | |
| if let some n := expr.constName? then | |
| let mod? ← ci.runMetaM i.lctx <| findModuleOf? n | |
| let modUri? ← match mod? with | |
| | some modName => | |
| let modFname? ← st.srcSearchPath.findWithExt ".lean" modName | |
| pure <| modFname?.map toFileUri | |
| | none => pure <| some doc.meta.uri | |
| let ranges? ← ci.runMetaM i.lctx <| findDeclarationRanges? n | |
| if let (some ranges, some modUri) := (ranges?, modUri?) then | |
| let declRangeToLspRange (r : DeclarationRange) : Lsp.Range := | |
| { start := ⟨r.pos.line - 1, r.charUtf16⟩ | |
| «end» := ⟨r.endPos.line - 1, r.endCharUtf16⟩ } | |
| let ll : LocationLink := { | |
| originSelectionRange? := some ⟨text.utf8PosToLspPos i.stx.getPos?.get!, | |
| text.utf8PosToLspPos i.stx.getTailPos?.get!⟩ | |
| targetUri := modUri | |
| targetRange := declRangeToLspRange ranges.range | |
| targetSelectionRange := declRangeToLspRange ranges.selectionRange | |
| } | |
| return #[ll] | |
| return #[] | |
| open Elab in | |
| partial def handlePlainGoal (p : PlainGoalParams) | |
| : ServerM (Task (Except IO.Error (Except RequestError (Option PlainGoal)))) := do | |
| let st ← read | |
| let doc ← st.docRef.get | |
| let text := doc.meta.text | |
| let hoverPos := text.lspPosToUtf8Pos p.position | |
| withWaitFindSnap doc (fun s => s.endPos > hoverPos) | |
| (notFoundX := return none) fun snap => do | |
| for t in snap.cmdState.infoState.trees do | |
| if let rs@(_ :: _) := t.goalsAt? hoverPos then | |
| let goals ← List.join <$> rs.mapM fun { ctxInfo := ci, tacticInfo := ti, useAfter := useAfter } => | |
| let ci := if useAfter then { ci with mctx := ti.mctxAfter } else { ci with mctx := ti.mctxBefore } | |
| let goals := if useAfter then ti.goalsAfter else ti.goalsBefore | |
| ci.runMetaM {} <| goals.mapM (fun g => Meta.withPPInaccessibleNames (Meta.ppGoal g)) | |
| let md := | |
| if goals.isEmpty then | |
| "no goals" | |
| else | |
| let goals := goals.map fun goal => s!"```lean | |
| {goal} | |
| ```" | |
| String.intercalate "\n---\n" goals | |
| return some { goals := goals.map toString |>.toArray, rendered := md } | |
| return none | |
| def hasRange (stx : Syntax) : Bool := | |
| stx.getPos?.isSome && stx.getTailPos?.isSome | |
| def rangeOfSyntax (text : FileMap) (stx : Syntax) : Range := | |
| ⟨text.utf8PosToLspPos <| stx.getPos?.get!, | |
| text.utf8PosToLspPos <| stx.getTailPos?.get!⟩ | |
| partial def handleDocumentHighlight (p : DocumentHighlightParams) : | |
| ServerM (Task (Except IO.Error (Except RequestError (Array DocumentHighlight)))) := do | |
| let doc ← (← read).docRef.get | |
| let text := doc.meta.text | |
| let pos := text.lspPosToUtf8Pos p.position | |
| let rec highlightReturn? (doRange? : Option Range) : Syntax → Option DocumentHighlight | |
| | stx@`(doElem|return%$i $e) => | |
| if stx.getPos?.get! <= pos && pos < stx.getTailPos?.get! then | |
| some { range := doRange?.getD (rangeOfSyntax text i), kind? := DocumentHighlightKind.text } | |
| else | |
| highlightReturn? doRange? e | |
| | `(do%$i $elems) => highlightReturn? (rangeOfSyntax text i) elems | |
| | stx => stx.getArgs.findSome? (highlightReturn? doRange?) | |
| withWaitFindSnap doc (fun s => s.endPos > pos) | |
| (notFoundX := pure #[]) fun snap => do | |
| if let some hi := highlightReturn? none snap.stx then | |
| return #[hi] | |
| return #[] | |
| partial def handleDocumentSymbol (p : DocumentSymbolParams) : | |
| ServerM (Task (Except IO.Error (Except RequestError DocumentSymbolResult))) := do | |
| let st ← read | |
| asTask do | |
| let doc ← st.docRef.get | |
| let ⟨cmdSnaps, e?⟩ ← doc.cmdSnaps.updateFinishedPrefix | |
| let mut stxs := cmdSnaps.finishedPrefix.map Snapshot.stx | |
| match e? with | |
| | some ElabTaskError.aborted => | |
| return Except.error RequestError.fileChanged | |
| | some _ => pure () -- TODO(WN): what to do on ioError? | |
| | none => | |
| let lastSnap := cmdSnaps.finishedPrefix.getLastD doc.headerSnap | |
| stxs := stxs ++ (← parseAhead doc.meta.text.source lastSnap).toList | |
| let (syms, _) := toDocumentSymbols doc.meta.text stxs | |
| return Except.ok { syms := syms.toArray } | |
| where | |
| toDocumentSymbols (text : FileMap) | |
| | [] => ([], []) | |
| | stx::stxs => match stx with | |
| | `(namespace $id) => sectionLikeToDocumentSymbols text stx stxs (id.getId.toString) SymbolKind.namespace id | |
| | `(section $(id)?) => sectionLikeToDocumentSymbols text stx stxs ((·.getId.toString) <$> id |>.getD "<section>") SymbolKind.namespace (id.getD stx) | |
| | `(end $(id)?) => ([], stx::stxs) | |
| | _ => | |
| let (syms, stxs') := toDocumentSymbols text stxs | |
| if stx.isOfKind ``Lean.Parser.Command.declaration && hasRange stx then | |
| let (name, selection) := match stx with | |
| | `($dm:declModifiers $ak:attrKind instance $[$np:namedPrio]? $[$id:ident$[.{$ls,*}]?]? $sig:declSig $val) => | |
| ((·.getId.toString) <$> id |>.getD s!"instance {sig.reprint.getD ""}", id.getD sig) | |
| | _ => match stx[1][1] with | |
| | `(declId|$id:ident$[.{$ls,*}]?) => (id.getId.toString, id) | |
| | _ => (stx[1][0].isIdOrAtom?.getD "<unknown>", stx[1][0]) | |
| if hasRange selection then | |
| (DocumentSymbol.mk { | |
| name := name | |
| kind := SymbolKind.method | |
| range := rangeOfSyntax text stx | |
| selectionRange := rangeOfSyntax text selection | |
| } :: syms, stxs') | |
| else | |
| (syms, stxs') | |
| else | |
| (syms, stxs') | |
| sectionLikeToDocumentSymbols (text : FileMap) (stx : Syntax) (stxs : List Syntax) (name : String) (kind : SymbolKind) (selection : Syntax) := | |
| let (syms, stxs') := toDocumentSymbols text stxs | |
| -- discard `end` | |
| let (syms', stxs'') := toDocumentSymbols text (stxs'.drop 1) | |
| let endStx := match stxs' with | |
| | endStx::_ => endStx | |
| | [] => (stx::stxs').getLast! | |
| (DocumentSymbol.mk { | |
| name := name | |
| kind := kind | |
| range := ⟨(rangeOfSyntax text stx).start, (rangeOfSyntax text endStx).«end»⟩ | |
| selectionRange := rangeOfSyntax text selection | |
| children? := syms.toArray | |
| } :: syms', stxs'') | |
| def noHighlightKinds : Array SyntaxNodeKind := #[ | |
| -- usually have special highlighting by the client | |
| ``Lean.Parser.Term.sorry, | |
| ``Lean.Parser.Term.type, | |
| ``Lean.Parser.Term.prop, | |
| -- not really keywords | |
| `antiquotName, | |
| ``Lean.Parser.Command.docComment] | |
| structure SemanticTokensContext where | |
| beginPos : String.Pos | |
| endPos : String.Pos | |
| text : FileMap | |
| infoState : Elab.InfoState | |
| structure SemanticTokensState where | |
| data : Array Nat | |
| lastLspPos : Lsp.Position | |
| partial def handleSemanticTokens (beginPos endPos : String.Pos) : | |
| ServerM (Task (Except IO.Error (Except RequestError SemanticTokens))) := do | |
| let st ← read | |
| let doc ← st.docRef.get | |
| let text := doc.meta.text | |
| let t ← doc.cmdSnaps.waitAll (·.beginPos < endPos) | |
| IO.mapTask (t := t) fun (snaps, _) => | |
| 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.cmdState.infoState) <| | |
| go s.stx | |
| return Except.ok { data := (← get).data } | |
| where | |
| go (stx : Syntax) := do | |
| match stx with | |
| | `($e.$id:ident) => go e; addToken id SemanticTokenType.property | |
| -- indistinguishable from next pattern | |
| --| `(level|$id:ident) => addToken id SemanticTokenType.variable | |
| | `($id:ident) => highlightId id | |
| | _ => | |
| if !noHighlightKinds.contains stx.getKind then | |
| highlightKeyword stx | |
| if stx.isOfKind choiceKind then | |
| go stx[0] | |
| else | |
| stx.getArgs.forM go | |
| highlightId (stx : Syntax) : ReaderT SemanticTokensContext (StateT SemanticTokensState IO) _ := do | |
| if let (some pos, some tailPos) := (stx.getPos?, stx.getTailPos?) then | |
| for t in (← read).infoState.trees do | |
| for ti in t.deepestNodes (fun | |
| | _, i@(Elab.Info.ofTermInfo ti) => match i.pos? with | |
| | some ipos => if pos <= ipos && ipos < tailPos then some ti else none | |
| | _ => none | |
| | _, _ => none) do | |
| match ti.expr with | |
| | Expr.fvar .. => addToken ti.stx SemanticTokenType.variable | |
| | _ => if ti.stx.getPos?.get! > pos then addToken ti.stx SemanticTokenType.property | |
| highlightKeyword stx := do | |
| if let Syntax.atom info val := stx then | |
| if val.bsize > 0 && val[0].isAlpha then | |
| addToken stx SemanticTokenType.keyword | |
| addToken stx type := do | |
| let ⟨beginPos, endPos, text, _⟩ ← read | |
| if let (some pos, some tailPos) := (stx.getPos?, stx.getTailPos?) then | |
| if beginPos <= pos && pos < endPos then | |
| let lspPos := (← get).lastLspPos | |
| let lspPos' := text.utf8PosToLspPos pos | |
| let deltaLine := lspPos'.line - lspPos.line | |
| let deltaStart := lspPos'.character - (if lspPos'.line == lspPos.line then lspPos.character else 0) | |
| let length := (text.utf8PosToLspPos tailPos).character - lspPos'.character | |
| let tokenType := type.toNat | |
| let tokenModifiers := 0 | |
| modify fun st => { | |
| data := st.data ++ #[deltaLine, deltaStart, length, tokenType, tokenModifiers] | |
| lastLspPos := lspPos' | |
| } | |
| def handleSemanticTokensFull (p : SemanticTokensParams) : | |
| ServerM (Task (Except IO.Error (Except RequestError SemanticTokens))) := do | |
| handleSemanticTokens 0 (1 <<< 16) | |
| def handleSemanticTokensRange (p : SemanticTokensRangeParams) : | |
| ServerM (Task (Except IO.Error (Except RequestError SemanticTokens))) := do | |
| let st ← read | |
| let doc ← st.docRef.get | |
| let text := doc.meta.text | |
| let beginPos := text.lspPosToUtf8Pos p.range.start | |
| let endPos := text.lspPosToUtf8Pos p.range.end | |
| handleSemanticTokens beginPos endPos | |
| partial def handleWaitForDiagnostics (p : WaitForDiagnosticsParams) | |
| : ServerM (Task (Except IO.Error (Except RequestError WaitForDiagnostics))) := do | |
| let st ← read | |
| let rec waitLoop : IO EditableDocument := do | |
| let doc ← st.docRef.get | |
| if p.version ≤ doc.meta.version then | |
| return doc | |
| else | |
| IO.sleep 50 | |
| waitLoop | |
| let t ← IO.asTask waitLoop | |
| let t ← IO.bindTask t fun | |
| | Except.error e => unreachable! | |
| | Except.ok doc => do | |
| let t₁ ← doc.cmdSnaps.waitAll | |
| return t₁.map fun _ => Except.ok WaitForDiagnostics.mk | |
| return t.map fun _ => Except.ok <| Except.ok WaitForDiagnostics.mk | |
| end RequestHandling | |
| section MessageHandling | |
| def parseParams (paramType : Type) [FromJson paramType] (params : Json) : ServerM paramType := | |
| match fromJson? params with | |
| | some parsed => pure parsed | |
| | none => throwServerError s!"Got param with wrong structure: {params.compress}" | |
| def handleNotification (method : String) (params : Json) : ServerM Unit := do | |
| let handle := fun paramType [FromJson paramType] (handler : paramType → ServerM Unit) => | |
| parseParams paramType params >>= handler | |
| match method with | |
| | "textDocument/didChange" => handle DidChangeTextDocumentParams handleDidChange | |
| | "$/cancelRequest" => handle CancelParams handleCancelRequest | |
| | _ => throwServerError s!"Got unsupported notification method: {method}" | |
| def queueRequest (id : RequestID) (requestTask : Task (Except IO.Error Unit)) | |
| : ServerM Unit := do | |
| updatePendingRequests (fun pendingRequests => pendingRequests.insert id requestTask) | |
| def handleRequest (id : RequestID) (method : String) (params : Json) | |
| : ServerM Unit := do | |
| let handle := fun paramType [FromJson paramType] respType [ToJson respType] | |
| (handler : paramType → RequestM respType) => do | |
| let st ← read | |
| let p ← parseParams paramType params | |
| let t ← handler p | |
| let t₁ ← (IO.mapTask · t) fun | |
| | Except.ok (Except.ok resp) => | |
| st.hOut.writeLspResponse ⟨id, resp⟩ | |
| | Except.ok (Except.error e) => | |
| st.hOut.writeLspResponseError { id := id, code := e.code, message := e.message } | |
| | Except.error e => | |
| st.hOut.writeLspResponseError { id := id, code := ErrorCode.internalError, message := toString e } | |
| queueRequest id t₁ | |
| match method with | |
| | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParams WaitForDiagnostics handleWaitForDiagnostics | |
| | "textDocument/completion" => handle CompletionParams CompletionList handleCompletion | |
| | "textDocument/hover" => handle HoverParams (Option Hover) handleHover | |
| | "textDocument/declaration" => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := false) | |
| | "textDocument/definition" => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := false) | |
| | "textDocument/typeDefinition" => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := true) | |
| | "textDocument/documentHighlight" => handle DocumentHighlightParams DocumentHighlightResult handleDocumentHighlight | |
| | "textDocument/documentSymbol" => handle DocumentSymbolParams DocumentSymbolResult handleDocumentSymbol | |
| | "textDocument/semanticTokens/full" => handle SemanticTokensParams SemanticTokens handleSemanticTokensFull | |
| | "textDocument/semanticTokens/range" => handle SemanticTokensRangeParams SemanticTokens handleSemanticTokensRange | |
| | "$/lean/plainGoal" => handle PlainGoalParams (Option PlainGoal) handlePlainGoal | |
| | _ => throwServerError s!"Got unsupported request: {method}" | |
| end MessageHandling | |
| section MainLoop | |
| partial def mainLoop : ServerM Unit := do | |
| let st ← read | |
| let msg ← st.hIn.readLspMessage | |
| let pendingRequests ← st.pendingRequestsRef.get | |
| let filterFinishedTasks (acc : PendingRequestMap) (id : RequestID) (task : Task (Except IO.Error Unit)) | |
| : ServerM PendingRequestMap := do | |
| if (←hasFinished task) then | |
| /- Handler tasks are constructed so that the only possible errors here | |
| are failures of writing a response into the stream. -/ | |
| if let Except.error e := task.get then | |
| throwServerError s!"Failed responding to request {id}: {e}" | |
| acc.erase id | |
| else acc | |
| let pendingRequests ← pendingRequests.foldM filterFinishedTasks pendingRequests | |
| st.pendingRequestsRef.set pendingRequests | |
| match msg with | |
| | Message.request id method (some params) => | |
| handleRequest id method (toJson params) | |
| mainLoop | |
| | Message.notification "exit" none => | |
| let doc ← st.docRef.get | |
| doc.cancelTk.set | |
| return () | |
| | Message.notification method (some params) => | |
| handleNotification method (toJson params) | |
| mainLoop | |
| | _ => throwServerError "Got invalid JSON-RPC message" | |
| end MainLoop | |
| def initAndRunWorker (i o e : FS.Stream) : IO UInt32 := do | |
| let i ← maybeTee "fwIn.txt" false i | |
| let o ← maybeTee "fwOut.txt" true o | |
| -- TODO(WN): act in accordance with InitializeParams | |
| let _ ← i.readLspRequestAs "initialize" InitializeParams | |
| let ⟨_, param⟩ ← i.readLspNotificationAs "textDocument/didOpen" DidOpenTextDocumentParams | |
| let doc := param.textDocument | |
| let meta : DocumentMeta := ⟨doc.uri, doc.version, doc.text.toFileMap⟩ | |
| let e ← e.withPrefix s!"[{param.textDocument.uri}] " | |
| let _ ← IO.setStderr e | |
| try | |
| let ctx ← initializeWorker meta i o e | |
| ReaderT.run (r := ctx) mainLoop | |
| return 0 | |
| catch e => | |
| IO.eprintln e | |
| publishDiagnostics meta #[{ range := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩, severity? := DiagnosticSeverity.error, message := e.toString }] o | |
| return 1 | |
| @[export lean_server_worker_main] | |
| def workerMain : IO UInt32 := do | |
| let i ← IO.getStdin | |
| let o ← IO.getStdout | |
| let e ← IO.getStderr | |
| try | |
| initAndRunWorker i o e | |
| catch err => | |
| e.putStrLn s!"worker initialization error: {err}" | |
| return (1 : UInt32) | |
| end Lean.Server.FileWorker |