chore: move RpcSession

src/Lean/Server/FileWorker.lean

@@ -18,6 +18,7 @@ import Lean.Server.AsyncList
 
 import Lean.Server.FileWorker.Utils
 import Lean.Server.FileWorker.RequestHandling
+import Lean.Server.Rpc.Basic
 
 /-!
 For general server architecture, see `README.md`. For details of IPC communication, see `Watchdog.lean`.

src/Lean/Server/FileWorker/Utils.lean

@@ -54,56 +54,4 @@ structure EditableDocument where
   cancelTk   : CancelToken
   deriving Inhabited
 
-/-- Concurrently modifiable part of an RPC session. -/
-structure RpcSessionState where
-  /-- Objects that are being kept alive for the RPC client, together with their type names,
-  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 : Std.PersistentHashMap Lsp.RpcRef (Name × NonScalar)
-  /-- Value to use for the next `RpcRef`. It is monotonically increasing to avoid any possible
-  bugs resulting from its reuse. -/
-  nextRef   : USize
-
-namespace RpcSessionState
-
-def store (st : RpcSessionState) (typeName : Name) (obj : NonScalar) : Lsp.RpcRef × RpcSessionState :=
-  let ref := ⟨st.nextRef⟩
-  let st' := { st with aliveRefs := st.aliveRefs.insert ref (typeName, obj)
-                       nextRef := st.nextRef + 1 }
-  (ref, st')
-
-def release (st : RpcSessionState) (ref : Lsp.RpcRef) : Bool × RpcSessionState :=
-  let released := st.aliveRefs.contains ref
-  (released, { st with aliveRefs := st.aliveRefs.erase ref })
-
-end RpcSessionState
-
-structure RpcSession where
-  sessionId       : UInt64
-  clientConnected : Bool
-  /-- We allow asynchronous elab tasks and request handlers to modify the state.
-  A single `Ref` ensures atomic transactions. -/
-  state           : IO.Ref RpcSessionState
-
-namespace RpcSession
-
-def new (clientConnected := false) : IO RpcSession := do
-  /- We generate a random ID to ensure that session IDs do not repeat across re-initializations
-  and worker restarts. Otherwise, the client may attempt to use outdated references. -/
-  let newId ← ByteArray.toUInt64LE! <$> IO.getRandomBytes 8
-  let newState ← IO.mkRef {
-    aliveRefs := Std.PersistentHashMap.empty
-    nextRef := 0
-  }
-  return {
-    sessionId := newId
-    clientConnected
-    state := newState
-  }
-
-end RpcSession
-
 end Lean.Server.FileWorker

src/Lean/Server/Requests.lean

@@ -55,7 +55,7 @@ def parseRequestParams (paramType : Type) [FromJson paramType] (params : Json)
       message := s!"Cannot parse request params: {params.compress}\n{inner}" }
 
 structure RequestContext where
-  rpcSesh       : FileWorker.RpcSession
+  rpcSesh       : RpcSession
   srcSearchPath : SearchPath
   doc           : FileWorker.EditableDocument
   hLog          : IO.FS.Stream
@@ -67,16 +67,9 @@ abbrev RequestM := ReaderT RequestContext <| ExceptT RequestError IO
 instance : Inhabited (RequestM α) :=
   ⟨throwThe IO.Error "executing Inhabited instance?!"⟩
 
-instance : MonadRpcSession RequestM where
-  rpcSessionId := do
-    (←read).rpcSesh.sessionId
-  rpcStoreRef typeName obj := do
-    (←read).rpcSesh.state.modifyGet fun st => st.store typeName obj
-  rpcGetRef r := do
-    let rs ← (←read).rpcSesh.state.get
-    rs.aliveRefs.find? r
-  rpcReleaseRef r := do
-    (←read).rpcSesh.state.modifyGet fun st => st.release r
+-- NOTE(WN): low priority so the entire `RequestContext` is preferred
+instance (priority := low) : MonadReaderOf RpcSession RequestM where
+  read := RequestContext.rpcSesh <$> read
 
 namespace RequestM
 open FileWorker

src/Lean/Server/Rpc/Basic.lean

@@ -31,6 +31,69 @@ instance {m n : Type → Type} [MonadLift m n] [MonadRpcSession m] : MonadRpcSes
   rpcGetRef r              := liftM (rpcGetRef r : m _)
   rpcReleaseRef r          := liftM (rpcReleaseRef r : m _)
 
+/-- Concurrently modifiable part of an RPC session. -/
+structure RpcSessionState where
+  /-- Objects that are being kept alive for the RPC client, together with their type names,
+  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 : Std.PersistentHashMap Lsp.RpcRef (Name × NonScalar)
+  /-- Value to use for the next `RpcRef`. It is monotonically increasing to avoid any possible
+  bugs resulting from its reuse. -/
+  nextRef   : USize
+
+namespace RpcSessionState
+
+def store (st : RpcSessionState) (typeName : Name) (obj : NonScalar) : Lsp.RpcRef × RpcSessionState :=
+  let ref := ⟨st.nextRef⟩
+  let st' := { st with aliveRefs := st.aliveRefs.insert ref (typeName, obj)
+                       nextRef := st.nextRef + 1 }
+  (ref, st')
+
+def release (st : RpcSessionState) (ref : Lsp.RpcRef) : Bool × RpcSessionState :=
+  let released := st.aliveRefs.contains ref
+  (released, { st with aliveRefs := st.aliveRefs.erase ref })
+
+end RpcSessionState
+
+structure RpcSession where
+  sessionId       : UInt64
+  clientConnected : Bool
+  /-- We allow asynchronous elab tasks and request handlers to modify the state.
+  A single `Ref` ensures atomic transactions. -/
+  state           : IO.Ref RpcSessionState
+
+namespace RpcSession
+
+def new (clientConnected := false) : IO RpcSession := do
+  /- We generate a random ID to ensure that session IDs do not repeat across re-initializations
+  and worker restarts. Otherwise, the client may attempt to use outdated references. -/
+  let newId ← ByteArray.toUInt64LE! <$> IO.getRandomBytes 8
+  let newState ← IO.mkRef {
+    aliveRefs := Std.PersistentHashMap.empty
+    nextRef := 0
+  }
+  return {
+    sessionId := newId
+    clientConnected
+    state := newState
+  }
+
+end RpcSession
+
+instance [Monad m] [MonadLiftT IO m] [MonadReaderOf RpcSession m] : MonadRpcSession m where
+  rpcSessionId := do
+    (←read).sessionId
+  rpcStoreRef typeName obj := do
+    liftM (m := IO) <| (←read).state.modifyGet fun st => st.store typeName obj
+  rpcGetRef r := do
+    let rs ← liftM (m := IO) <| (←read).state.get
+    rs.aliveRefs.find? r
+  rpcReleaseRef r := do
+    liftM (m := IO) <| (←read).state.modifyGet fun st => st.release r
+
 /-- `RpcEncoding α β` means that `α` may participate in RPC calls with its on-the-wire LSP encoding
 being `β`. This is useful when `α` contains fields which must be marshalled in a special way. In
 particular, we encode `WithRpcRef` fields as opaque references rather than send their content.