feat: initial RPC

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

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

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 _)

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

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

src/Lean/Server/FileWorker/Utils.lean

@@ -53,4 +53,12 @@ structure EditableDocument where
   cancelTk   : CancelToken
   deriving Inhabited
 
-end Lean.Server.FileWorker
+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

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
+end Lean.Server.Requests