From 3ec568c11071b28ed5499b4ac2fb6d52b73b3f99 Mon Sep 17 00:00:00 2001 From: Wojciech Nawrocki Date: Sun, 11 Jul 2021 22:36:11 -0700 Subject: [PATCH] feat: initial RPC --- src/Lean/Data/Lsp/Extra.lean | 34 +++++- src/Lean/Server/FileSource.lean | 43 +++---- src/Lean/Server/FileWorker.lean | 7 +- .../Server/FileWorker/RequestHandling.lean | 33 ++++++ src/Lean/Server/FileWorker/Rpc.lean | 111 ++++++++++++++++++ src/Lean/Server/FileWorker/Utils.lean | 10 +- src/Lean/Server/Requests.lean | 3 +- 7 files changed, 215 insertions(+), 26 deletions(-) create mode 100644 src/Lean/Server/FileWorker/Rpc.lean diff --git a/src/Lean/Data/Lsp/Extra.lean b/src/Lean/Data/Lsp/Extra.lean index 2237e2c71d..30da32f062 100644 --- a/src/Lean/Data/Lsp/Extra.lean +++ b/src/Lean/Data/Lsp/Extra.lean @@ -20,7 +20,9 @@ The following additional requests and notifications are supported: - `$/lean/fileProgress`: Notification that contains the ranges of the document that are still being processed by the server. -- `$/lean/plainGoal`: Returns the goal(s) at the specified position, pretty-printed as a string. +- `$/lean/plainGoal`/`$/lean/plainTermGoal`: Returns the goal(s) at the specified position, pretty-printed as a string. + +- `$/lean/rpc/...`: See `Lean.Server.FileWorker.Rpc`. -/ namespace Lean.Lsp @@ -30,7 +32,7 @@ open Json structure WaitForDiagnosticsParams where uri : DocumentUri version : Nat - deriving ToJson, FromJson + deriving FromJson, ToJson structure WaitForDiagnostics @@ -65,4 +67,32 @@ structure PlainTermGoal where range : Range deriving FromJson, ToJson +/-- An object which RPC clients can refer to without marshalling. -/ +abbrev RpcRef := USize +inductive RpcValue where + | json : Json → RpcValue + | ref : RpcRef → RpcValue + deriving FromJson, ToJson, Inhabited + +inductive RpcValueKind where + | json | ref + deriving FromJson, ToJson + +/-- A request for Lean to execute a function previously bound for RPC. + +Extending TDPP is weird. But in Lean, symbols exist in the context of a position +within a source file. So we need this to call user code in the env at that position. -/ +structure RpcCallParams extends TextDocumentPositionParams where + sessionId : USize + /-- Function to invoke. Must be fully qualified. -/ + method : Name + args : Array RpcValue + /-- How the return value should be sent. -/ + retKind : RpcValueKind + deriving FromJson, ToJson + +structure RpcCallResult where + value : RpcValue + deriving FromJson, ToJson + end Lean.Lsp diff --git a/src/Lean/Server/FileSource.lean b/src/Lean/Server/FileSource.lean index 362f7e50d5..a762fc0ec3 100644 --- a/src/Lean/Server/FileSource.lean +++ b/src/Lean/Server/FileSource.lean @@ -12,67 +12,70 @@ class FileSource (α : Type) where fileSource : α → DocumentUri export FileSource (fileSource) -instance Location.hasFileSource : FileSource Location := +instance : FileSource Location := ⟨fun l => l.uri⟩ -instance TextDocumentIdentifier.hasFileSource : FileSource TextDocumentIdentifier := +instance : FileSource TextDocumentIdentifier := ⟨fun i => i.uri⟩ -instance VersionedTextDocumentIdentifier.hasFileSource : FileSource VersionedTextDocumentIdentifier := +instance : FileSource VersionedTextDocumentIdentifier := ⟨fun i => i.uri⟩ -instance TextDocumentEdit.hasFileSource : FileSource TextDocumentEdit := +instance : FileSource TextDocumentEdit := ⟨fun e => fileSource e.textDocument⟩ -instance TextDocumentItem.hasFileSource : FileSource TextDocumentItem := +instance : FileSource TextDocumentItem := ⟨fun i => i.uri⟩ -instance TextDocumentPositionParams.hasFileSource : FileSource TextDocumentPositionParams := +instance : FileSource TextDocumentPositionParams := ⟨fun p => fileSource p.textDocument⟩ -instance DidOpenTextDocumentParams.hasFileSource : FileSource DidOpenTextDocumentParams := +instance : FileSource DidOpenTextDocumentParams := ⟨fun p => fileSource p.textDocument⟩ -instance DidChangeTextDocumentParams.hasFileSource : FileSource DidChangeTextDocumentParams := +instance : FileSource DidChangeTextDocumentParams := ⟨fun p => fileSource p.textDocument⟩ -instance DidCloseTextDocumentParams.hasFileSource : FileSource DidCloseTextDocumentParams := +instance : FileSource DidCloseTextDocumentParams := ⟨fun p => fileSource p.textDocument⟩ -instance CompletionParams.hasFileSource : FileSource CompletionParams := +instance : FileSource CompletionParams := ⟨fun h => fileSource h.toTextDocumentPositionParams⟩ -instance HoverParams.hasFileSource : FileSource HoverParams := +instance : FileSource HoverParams := ⟨fun h => fileSource h.toTextDocumentPositionParams⟩ -instance DeclarationParams.hasFileSource : FileSource DeclarationParams := +instance : FileSource DeclarationParams := ⟨fun h => fileSource h.toTextDocumentPositionParams⟩ -instance DefinitionParams.hasFileSource : FileSource DefinitionParams := +instance : FileSource DefinitionParams := ⟨fun h => fileSource h.toTextDocumentPositionParams⟩ -instance TypeDefinitionParams.hasFileSource : FileSource TypeDefinitionParams := +instance : FileSource TypeDefinitionParams := ⟨fun h => fileSource h.toTextDocumentPositionParams⟩ -instance WaitForDiagnosticsParams.hasFileSource : FileSource WaitForDiagnosticsParams := +instance : FileSource WaitForDiagnosticsParams := ⟨fun p => p.uri⟩ -instance DocumentHighlightParams.hasFileSource : FileSource DocumentHighlightParams := +instance : FileSource DocumentHighlightParams := ⟨fun h => fileSource h.toTextDocumentPositionParams⟩ -instance DocumentSymbolParams.hasFileSource : FileSource DocumentSymbolParams := +instance : FileSource DocumentSymbolParams := ⟨fun p => fileSource p.textDocument⟩ -instance SemanticTokensParams.hasFileSource : FileSource SemanticTokensParams := +instance : FileSource SemanticTokensParams := ⟨fun p => fileSource p.textDocument⟩ -instance SemanticTokensRangeParams.hasFileSource : FileSource SemanticTokensRangeParams := +instance : FileSource SemanticTokensRangeParams := ⟨fun p => fileSource p.textDocument⟩ -instance PlainGoalParams.hasFileSource : FileSource PlainGoalParams := +instance : FileSource PlainGoalParams := ⟨fun p => fileSource p.textDocument⟩ instance : FileSource PlainTermGoalParams where fileSource p := fileSource p.textDocument +instance : FileSource RpcCallParams where + fileSource p := fileSource p.textDocument + end Lean.Lsp diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean index a3caa6ef1c..d9d3d9eea7 100644 --- a/src/Lean/Server/FileWorker.lean +++ b/src/Lean/Server/FileWorker.lean @@ -97,8 +97,9 @@ structure WorkerContext where hOut : FS.Stream hLog : FS.Stream srcSearchPath : SearchPath - docRef : IO.Ref EditableDocument + docRef : IO.Ref EditableDocument -- TODO WorkerM := StateT pendingRequestsRef : IO.Ref PendingRequestMap + rpcSeshRef : IO.Ref (Option RpcSession) abbrev WorkerM := ReaderT WorkerContext IO @@ -192,6 +193,7 @@ section Initialization srcSearchPath := srcSearchPath docRef := ←IO.mkRef doc pendingRequestsRef := ←IO.mkRef RBMap.empty + rpcSeshRef := ←IO.mkRef none } end Initialization @@ -289,7 +291,8 @@ section MessageHandling : WorkerM Unit := do let st ← read let rc : Requests.RequestContext := - { srcSearchPath := st.srcSearchPath + { rpcSesh := ← st.rpcSeshRef.get + srcSearchPath := st.srcSearchPath docRef := st.docRef hLog := st.hLog } let t? ← (ExceptT.run <| Requests.handleLspRequest method params rc : IO _) diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean index 1135cceca4..661adf0444 100644 --- a/src/Lean/Server/FileWorker/RequestHandling.lean +++ b/src/Lean/Server/FileWorker/RequestHandling.lean @@ -10,6 +10,7 @@ import Lean.Data.Json import Lean.Data.Lsp import Lean.Server.FileWorker.Utils +import Lean.Server.FileWorker.Rpc import Lean.Server.Requests import Lean.Server.Completion @@ -369,6 +370,37 @@ partial def handleWaitForDiagnostics (p : WaitForDiagnosticsParams) let t₁ ← doc.cmdSnaps.waitAll return t₁.map fun _ => pure WaitForDiagnostics.mk +def handleRpcCall (p : RpcCallParams) : RequestM (RequestTask RpcCallResult) := do + let rc ← read + let some rpcSesh ← rc.rpcSesh + | throwThe RequestError { code := JsonRpc.ErrorCode.unknownErrorCode -- TODO + message := s!"RPC session not initialized" } + if p.sessionId ≠ rpcSesh.sessionId then + throwThe RequestError { code := JsonRpc.ErrorCode.unknownErrorCode -- TODO + message := s!"Outdated RPC session" } + let some proc ← (← Rpc.rpcProcedures.get).find? p.method + | throwThe RequestError { code := JsonRpc.ErrorCode.methodNotFound + message := s!"No RPC method '{p.method}' bound" } + + let aliveRefs ← rpcSesh.aliveRefs.get + let args ← p.args.mapM fun + | RpcValue.json j => pure <| Rpc.RpcInput.json j + | RpcValue.ref p => + match aliveRefs.find? p with + | none => throwThe RequestError { code := JsonRpc.ErrorCode.invalidParams + message := s!"reference '{p}' is not valid" } + | some r => pure <| Rpc.RpcInput.ref r + let ret ← proc.wrapper args p.retKind + + let ret' ← match ret with + | Rpc.RpcInput.json j => pure <| RpcValue.json j + | Rpc.RpcInput.ref r => + let uid := r.ptr -- TODO + rpcSesh.aliveRefs.modify fun refs => refs.insert uid r + pure <| RpcValue.ref uid + + RequestM.asTask (return ⟨ ret' ⟩) -- TODO Task + builtin_initialize registerLspRequestHandler "textDocument/waitForDiagnostics" WaitForDiagnosticsParams WaitForDiagnostics handleWaitForDiagnostics registerLspRequestHandler "textDocument/completion" CompletionParams CompletionList handleCompletion @@ -382,5 +414,6 @@ builtin_initialize registerLspRequestHandler "textDocument/semanticTokens/range" SemanticTokensRangeParams SemanticTokens handleSemanticTokensRange registerLspRequestHandler "$/lean/plainGoal" PlainGoalParams (Option PlainGoal) handlePlainGoal registerLspRequestHandler "$/lean/plainTermGoal" PlainTermGoalParams (Option PlainTermGoal) handlePlainTermGoal + registerLspRequestHandler "$/lean/rpc/call" RpcCallParams RpcCallResult handleRpcCall end Lean.Server.FileWorker diff --git a/src/Lean/Server/FileWorker/Rpc.lean b/src/Lean/Server/FileWorker/Rpc.lean new file mode 100644 index 0000000000..fe62f6da2d --- /dev/null +++ b/src/Lean/Server/FileWorker/Rpc.lean @@ -0,0 +1,111 @@ +import Lean.Elab.Command +import Lean.Meta.Basic +import Lean.Data.Lsp.Extra + +/-! Allows LSP clients to make Remote Procedure Calls to the server. + +The single use case for these is to allow the infoview UI to refer to and manipulate heavy-weight +objects residing on the server. It would be inefficient to serialize these into JSON and send over. +For example, the client can format an `Expr` without transporting the whole `Environment`. +-/ + +namespace Lean.Server.Rpc +open Lsp (RpcValueKind) + +/-- Like `RpcValue`, but with decoded `ref` IDs. -/ +inductive RpcInput where + | ref : IO.UntypedRef → RpcInput + | json : Json → RpcInput + deriving Inhabited +open RpcInput + +private unsafe def RpcSession.decodePtrOrJson [Inhabited α] [FromJson α] : RpcInput → IO α + | ref r => r.getAsUnsafe α + | json j => match @fromJson? α _ j with + | Except.error e => throwThe IO.Error e + | Except.ok v => v + +private unsafe def RpcSession.decodePtr [Inhabited α] : RpcInput → IO α + | ref r => r.getAsUnsafe α + | json j => throwThe IO.Error "cannot decode, no FromJson instance" + +open Elab Meta in +elab "wrap_rpc_func" funcId:ident : command => do + let env ← getEnv + let funcNm := funcId.getId + match env.find? funcNm with + | none => throwError "unknown constant '{funcNm}'" + | some info => + let cmds ← Command.liftTermElabM none do + forallTelescopeReducing info.type fun args retT => do + let mkFreshFVarName (arg : Expr) : TermElabM Name := do + let localDecl ← getLocalDecl arg.fvarId! + mkFreshUserName localDecl.userName.eraseMacroScopes + + let hasInstance (cl : Name) (tp : Expr) : TermElabM Bool := do + let clTp ← + try mkAppM cl #[tp] + catch _ => throwError "failed to construct '{cl} {tp}'" + let inst? ← trySynthInstance clTp + inst? matches LOption.some _ + + -- (args : Array RpcInput) (retKind : RpcValueKind) + let rpcArgsId ← mkIdent <$> mkFreshUserName `args + let retKindArgId ← mkIdent <$> mkFreshUserName `retKind + + let mut argTNames := #[] + let mut letIds := #[] + let mut decoders := #[] + let mut i := 0 -- enumerate? :( + for arg in args do + let argName ← mkFreshFVarName arg + let argT ← inferType arg + let some argTName ← argT.constName? + | throwError "non-constant type '{argT}'" + argTNames := argTNames.push <| mkIdent argTName + letIds := letIds.push <| mkIdent argName + -- a term to decode the `i`th argument + -- TODO add [FromJson α] for polymorphic arguments + let decoder ← + if (← hasInstance ``FromJson argT) then + `(Parser.Term.doExpr| RpcSession.decodePtrOrJson $rpcArgsId:ident[$(quote i)] ) + else + `(Parser.Term.doExpr| RpcSession.decodePtr $rpcArgsId:ident[$(quote i)] ) + decoders := decoders.push decoder + i := i+1 + + let rpcUnsafeId := mkIdent <| Name.mkStr funcNm "__rpc_unsafe" + let mkRetJson ← + if (← hasInstance ``ToJson retT) then + `(Parser.Term.doSeqIndent| pure <| RpcInput.json <| toJson <| @$funcId:ident $letIds*) + else + `(Parser.Term.doSeqIndent| throwThe IO.Error "cannot encode, no ToJson instance") + let cmd₁ ← `( + private unsafe def $rpcUnsafeId:ident ($rpcArgsId : Array RpcInput) ($retKindArgId : RpcValueKind) + : IO RpcInput := do + $[let $letIds:ident : $argTNames:ident ← $decoders]* + let ret ← match $retKindArgId:ident with + | RpcValueKind.ref => RpcInput.ref <$> IO.UntypedRef.mkRefUnsafe (@$funcId:ident $letIds*) + | RpcValueKind.json => $mkRetJson + return ret + ) + + let rpcId := mkIdent <| Name.mkStr funcNm "__rpc" + let cmd₂ ← `( + @[implementedBy $rpcUnsafeId:ident] + private constant $rpcId:ident ($rpcArgsId : Array RpcInput) ($retKindArgId : RpcValueKind) + : IO RpcInput + ) + + return #[cmd₁, cmd₂] + cmds.forM Command.elabCommand + +private structure RpcProcedure where + wrapper : Array RpcInput → RpcValueKind → IO RpcInput + +builtin_initialize rpcProcedures : IO.Ref (Std.PersistentHashMap Name RpcProcedure) ← + IO.mkRef {} + +-- TODO def registerRpcProcedure + +end Lean.Server.Rpc diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean index c72c123586..21241e46c9 100644 --- a/src/Lean/Server/FileWorker/Utils.lean +++ b/src/Lean/Server/FileWorker/Utils.lean @@ -53,4 +53,12 @@ structure EditableDocument where cancelTk : CancelToken deriving Inhabited -end Lean.Server.FileWorker \ No newline at end of file +structure RpcSession where + sessionId : USize + /-- Objects that are being kept alive for the RPC client, mapped to by their RPC reference. + Note that we may currently have multiple references to the same object. It is only disposed + of once all of those are gone. This simplifies the client a bit as it can drop every reference + received separately. -/ + aliveRefs : IO.Ref (Std.PersistentHashMap USize UntypedRef) + +end Lean.Server.FileWorker diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean index 845885afb8..3a7d7ae951 100644 --- a/src/Lean/Server/Requests.lean +++ b/src/Lean/Server/Requests.lean @@ -47,6 +47,7 @@ def toLspResponseError (id : RequestID) (e : RequestError) : ResponseError Unit end RequestError structure RequestContext where + rpcSesh : Option FileWorker.RpcSession srcSearchPath : SearchPath docRef : IO.Ref FileWorker.EditableDocument hLog : IO.FS.Stream @@ -165,4 +166,4 @@ def handleLspRequest (method : String) (params : Json) : RequestM (RequestTask J | some rh => rh.handle params end HandlerTable -end Lean.Server.Requests \ No newline at end of file +end Lean.Server.Requests