documentation & cleanup in FileWorker

server/GameServer/FileWorker.lean

@@ -5,6 +5,7 @@ import GameServer.ImportModules
 import GameServer.SaveData
 
 namespace MyModule
+
 open Lean
 open Elab
 open Parser
@@ -48,7 +49,8 @@ partial def parseTactic (inputCtx : InputContext) (pmctx : ParserModuleContext)
 
 end MyModule
 
-namespace MyServer.FileWorker
+namespace GameServer.FileWorker
+
 open Lean
 open Lean.Server
 open Lean.Server.FileWorker
@@ -57,22 +59,50 @@ open IO
 open Snapshots
 open JsonRpc
 
-structure GameWorkerState :=
+/--
+Game-specific state to be packed on top of the `Lean.Server.FileWorker.WorkerState`
+used by the lean server.
+-/
+structure WorkerState :=
+  /--
+  Collection of items which are considered unlocked.
+  Tactics and theorems are mixed together.
+  -/
   inventory : Array String
   /--
-  Check for tactics/theorems that are not unlocked.
+  Difficulty determines whether tactics/theorems can be locked.
-  0: no check
+  * 0: do not check
-  1: give warnings
+  * 1: give warnings when locked items are used
-  2: give errors
+  * 2: give errors when locked items are used
   -/
   difficulty : Nat
+  /--
+  `levelInfo` contains all the (static) information about the level which is not influenced
+  by the user's progress.
+  -/
   levelInfo : LevelInfo
 deriving ToJson, FromJson
 
-abbrev GameWorkerM := StateT GameWorkerState Server.FileWorker.WorkerM
+/--
+Pack the `GameServer.FileWorker.WorkerState` on top of the normal worker monad
+`Server.FileWorker.WorkerM`.
+-/
+abbrev WorkerM := StateT WorkerState Server.FileWorker.WorkerM
 
 section Elab
 
+/-- Add a message. use `(severity := .warning)` to specify the severity-/
+def addMessage (info : SourceInfo) (inputCtx : Parser.InputContext)
+    (severity := MessageSeverity.warning) (s : MessageData) :
+    Elab.Command.CommandElabM Unit := do
+  modify fun st => { st with
+    messages := st.messages.add {
+      fileName := inputCtx.fileName
+      severity := severity
+      pos      := inputCtx.fileMap.toPosition (info.getPos?.getD 0)
+      data     := s }}
+
+/-- Deprecated! -/
 def addErrorMessage (info : SourceInfo) (inputCtx : Parser.InputContext) (s : MessageData) :
     Elab.Command.CommandElabM Unit := do
   modify fun st => { st with
@@ -80,81 +110,95 @@ def addErrorMessage (info : SourceInfo) (inputCtx : Parser.InputContext) (s : Me
       fileName := inputCtx.fileName
       severity := MessageSeverity.error
       pos      := inputCtx.fileMap.toPosition (info.getPos?.getD 0)
-        data     := s
+      data     := s }}
-      }
-    }
 
 -- TODO: use HashSet for allowed tactics?
-/-- Find all tactics in syntax object that are forbidden according to a
+/--
-set `allowed` of allowed tactics. -/
+Find all tactics in syntax object that are forbidden according to a
-partial def findForbiddenTactics (inputCtx : Parser.InputContext)
+set `allowed` of allowed tactics.
-    (gameWorkerState : GameWorkerState) (stx : Syntax) :
+-/
-    Elab.Command.CommandElabM Unit := do
+partial def findForbiddenTactics (inputCtx : Parser.InputContext) (workerState : WorkerState)
-  let levelInfo := gameWorkerState.levelInfo
+    (stx : Syntax) : Elab.Command.CommandElabM Unit := do
+  let levelInfo := workerState.levelInfo
+  -- Parse the syntax object and look for tactics and declarations.
   match stx with
   | .missing => return ()
   | .node _info _kind args =>
+    -- Go inside a node.
     for arg in args do
-      findForbiddenTactics inputCtx gameWorkerState arg
+      findForbiddenTactics inputCtx workerState arg
   | .atom info val =>
-    -- Whitelisted keywords.
+    -- Atoms might be tactic names or other keywords.
-    -- Note: We need a whitelist because we cannot really distinguish keywords from tactics.
+    -- Note: We whitelisted known keywords because we cannot
+    -- distinguish keywords from tactic names.
     let allowed := ["with", "fun", "at", "only", "by", "to", "generalizing", "says"]
-    -- ignore syntax elements that do not start with a letter or are allowed
+    -- Ignore syntax elements that do not start with a letter or are listed above.
     if 0 < val.length ∧ val.data[0]!.isAlpha ∧ not (allowed.contains val) then
-      let val := val.dropRightWhile (fun c => c == '!' || c == '?') -- treat `simp?` and `simp!` like `simp`
+      -- Treat `simp?` and `simp!` like `simp`
+      let val := val.dropRightWhile (fun c => c == '!' || c == '?')
       match levelInfo.tactics.find? (·.name.toString == val) with
       | none =>
-        -- Note: This case means that the tactic will never be introduced in the game.
+        -- Tactic will never be introduced in the game.
-        match gameWorkerState.inventory.find? (· == val) with
+        match workerState.inventory.find? (· == val) with
+        | some _ =>
+          -- Tactic is in the inventory, allow it.
+          -- Note: This case shouldn't be possible...
+          pure ()
         | none =>
-          addWarningMessage info s!"You have not unlocked the tactic '{val}' yet!"
+          -- Tactic is not in the inventory.
-        | some _ => pure () -- tactic is in the inventory, allow it.
+          addMessageByDifficulty info s!"The tactic '{val}' is not available in this game!"
       | some tac =>
+        -- Tactic is introduced at some point in the game.
         if tac.locked then
-          match gameWorkerState.inventory.find? (· == val) with
+          match workerState.inventory.find? (· == val) with
           | none =>
-            addWarningMessage info s!"You have not unlocked the tactic '{val}' yet!"
+            -- Tactic is marked as locked and not in the inventory.
-          | some _ => pure () -- tactic is in the inventory, allow it.
+            addMessageByDifficulty info s!"You have not unlocked the tactic '{val}' yet!"
+          | some _ =>
+            -- Tactic is in the inventory, allow it.
+            pure ()
         else if tac.disabled then
-          addWarningMessage info s!"The tactic '{val}' is disabled in this level!"
+          -- Tactic is disabled in this level.
+          addMessageByDifficulty info s!"The tactic '{val}' is disabled in this level!"
   | .ident info _rawVal val _preresolved =>
+      -- Try to resolve the name
       let ns ←
         try resolveGlobalConst (mkIdent val)
-        catch | _ => pure [] -- catch "unknown constant" error
+        -- Catch "unknown constant" error
+        catch | _ => pure []
       for n in ns do
         let some (.thmInfo ..) := (← getEnv).find? n
-          | return () -- not a theorem -> ignore
+          -- Not a theorem, no checks needed.
-        -- Forbid the theorem we are proving currently
+          | return ()
         if some n = levelInfo.statementName then
-          addErrorMessage info inputCtx s!"Structural recursion: you can't use '{n}' to proof itself!"
+          -- Forbid the theorem we are proving currently
-
+          addMessage info inputCtx (severity := .error)
-        let lemmasAndDefs := levelInfo.lemmas ++ levelInfo.definitions
+            s!"Structural recursion: you can't use '{n}' to proof itself!"
-        match lemmasAndDefs.find? (fun l => l.name == n) with
+        let theoremsAndDefs := levelInfo.lemmas ++ levelInfo.definitions
-        | none => addWarningMessage info s!"You have not unlocked the lemma/definition '{n}' yet!"
+        match theoremsAndDefs.find? (·.name == n) with
-        | some lem =>
+        | none =>
-          if lem.locked then
+          -- Theorem will never be introduced in this game
-            addWarningMessage info s!"You have not unlocked the lemma/definition '{n}' yet!"
+          addMessageByDifficulty info s!"You have not unlocked the theorem/definition '{n}' yet!"
-          else if lem.disabled then
+        | some thm =>
-            addWarningMessage info s!"The lemma/definition '{n}' is disabled in this level!"
+          -- Theorem is introduced at some point in the game.
-where addWarningMessage (info : SourceInfo) (s : MessageData) :=
+          if thm.locked then
-  let difficulty := gameWorkerState.difficulty
+            -- Theorem is still locked.
+            addMessageByDifficulty info s!"You have not unlocked the theorem/definition '{n}' yet!"
+          else if thm.disabled then
+            -- Theorem is disabled in this level.
+            addMessageByDifficulty info s!"The theorem/definition '{n}' is disabled in this level!"
+where addMessageByDifficulty (info : SourceInfo) (s : MessageData) :=
+  -- See `GameServer.FileWorker.WorkerState.difficulty`. Send nothing/warnings/errors
+  -- deppending on difficulty.
+  let difficulty := workerState.difficulty
   if difficulty > 0 then
-    modify fun st => { st with
+    addMessage info inputCtx (if difficulty > 1 then .error else .warning) s
-      messages := st.messages.add {
-        fileName := inputCtx.fileName
-        severity := if difficulty > 1 then MessageSeverity.error else MessageSeverity.warning
-        pos      := inputCtx.fileMap.toPosition (info.getPos?.getD 0)
-        data     := s
-      }
-    }
   else pure ()
--- where addErrorMessage (info : SourceInfo) (s : MessageData) :=
---   pure ()
 
 open Elab Meta Expr in
+
 def compileProof (inputCtx : Parser.InputContext) (snap : Snapshot) (hasWidgets : Bool)
-    (couldBeEndSnap : Bool) (gameWorkerState : GameWorkerState)
+    (couldBeEndSnap : Bool) (gameWorkerState : WorkerState)
     (initParams : Lsp.InitializeParams) : IO Snapshot := do
   -- Recognize end snap
   if inputCtx.input.atEnd snap.mpState.pos ∧ couldBeEndSnap then
@@ -288,7 +332,7 @@ where
 
   /-- Elaborates the next command after `parentSnap` and emits diagnostics into `hOut`. -/
   private def nextSnap (ctx : WorkerContext) (m : DocumentMeta) (cancelTk : CancelToken)
-      (gameWorkerState : GameWorkerState) (initParams : Lsp.InitializeParams)
+      (gameWorkerState : WorkerState) (initParams : Lsp.InitializeParams)
       : AsyncElabM (Option Snapshot) := do
     cancelTk.check
     let s ← get
@@ -328,7 +372,7 @@ where
 
   /-- Elaborates all commands after the last snap (at least the header snap is assumed to exist), emitting the diagnostics into `hOut`. -/
   def unfoldSnaps (m : DocumentMeta) (snaps : Array Snapshot) (cancelTk : CancelToken)
-      (startAfterMs : UInt32) (gameWorkerState : GameWorkerState)
+      (startAfterMs : UInt32) (gameWorkerState : WorkerState)
       : ReaderT WorkerContext IO (AsyncList ElabTaskError Snapshot) := do
     let ctx ← read
     let some headerSnap := snaps[0]? | panic! "empty snapshots"
@@ -351,7 +395,7 @@ end Elab
 section Updates
 
 /-- Given the new document, updates editable doc state. -/
-  def updateDocument (newMeta : DocumentMeta) : GameWorkerM Unit := do
+def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := do
   let s ← get
   let ctx ← read
   let oldDoc := (← StateT.lift get).doc
@@ -469,7 +513,7 @@ section Initialization
     return (headerSnap, srcSearchPath)
 
 def initializeWorker (meta : DocumentMeta) (i o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
-      (gameDir : String) (gameWorkerState : GameWorkerState) : IO (WorkerContext × WorkerState) := do
+    (gameDir : String) (gameWorkerState : WorkerState) : IO (WorkerContext × Server.FileWorker.WorkerState) := do
   let clientHasWidgets := initParams.initializationOptions?.bind (·.hasWidgets?) |>.getD false
 
   let (headerStx, headerTask) ← compileHeader meta o opts (hasWidgets := clientHasWidgets)
@@ -500,7 +544,7 @@ end Initialization
 
 section NotificationHandling
 
-  def handleDidChange (p : DidChangeTextDocumentParams) : GameWorkerM Unit := do
+def handleDidChange (p : DidChangeTextDocumentParams) : WorkerM Unit := do
   let docId := p.textDocument
   let changes := p.contentChanges
   let oldDoc := (← StateT.lift get).doc
@@ -517,26 +561,39 @@ end NotificationHandling
 
 section MessageHandling
 
-  def handleNotification (method : String) (params : Json) : GameWorkerM Unit := do
+
-    let handle := fun paramType [FromJson paramType] (handler : paramType → GameWorkerM Unit) =>
+/--
+Modified notification handler.
+
+Compare to `Lean.Server.FileWorker.handleNotification`.
+We use the modified `WorkerM` and use our custom `handleDidChange`.
+
+-/
+def handleNotification (method : String) (params : Json) : WorkerM Unit := do
+  let handle := fun paramType [FromJson paramType] (handler : paramType → WorkerM Unit) =>
     (StateT.lift <| parseParams paramType params) >>= handler
   match method with
+  -- Modified `textDocument/didChange`, using a custom `handleDidChange`
   | "textDocument/didChange" => handle DidChangeTextDocumentParams (handleDidChange)
+  -- unmodified
   | "$/cancelRequest"        => handle CancelParams (handleCancelRequest ·)
-    | "$/setTrace"        => pure ()
+  -- unmodified
   | "$/lean/rpc/release"     => handle RpcReleaseParams (handleRpcRelease ·)
+  -- unmodified
   | "$/lean/rpc/keepAlive"   => handle RpcKeepAliveParams (handleRpcKeepAlive ·)
+  -- New. TODO: What is this for?
+  | "$/setTrace"        => pure ()
   | _                        => throwServerError s!"Got unsupported notification method: {method}"
 
 end MessageHandling
 
 section MainLoop
-  partial def mainLoop : GameWorkerM Unit := do
+
-    let ctx ← read
+/--
-    let mut st ← StateT.lift get
+Erase finished tasks if there are no errors.
-    let msg ← ctx.hIn.readLspMessage
+-/
-    let filterFinishedTasks (acc : PendingRequestMap) (id : RequestID) (task : Task (Except IO.Error Unit))
+private def filterFinishedTasks (acc : PendingRequestMap) (id : RequestID)
-        : IO PendingRequestMap := do
+    (task : Task (Except IO.Error Unit)) : IO 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. -/
@@ -544,9 +601,17 @@ section MainLoop
       throwServerError s!"Failed responding to request {id}: {e}"
     pure <| acc.erase id
   else pure acc
-    let pendingRequests ← st.pendingRequests.foldM (fun acc id task => filterFinishedTasks acc id task) st.pendingRequests
-    st := { st with pendingRequests }
 
+/--
+The main-loop.
+-/
+partial def mainLoop : WorkerM Unit := do
+  let ctx ← read
+  let mut st ← StateT.lift get
+  let msg ← ctx.hIn.readLspMessage
+  let pendingRequests ← st.pendingRequests.foldM (fun acc id task =>
+    filterFinishedTasks acc id task) st.pendingRequests
+  st := { st with pendingRequests }
   -- Opportunistically (i.e. when we wake up on messages) check if any RPC session has expired.
   for (id, seshRef) in st.rpcSessions do
     let sesh ← seshRef.get
@@ -554,23 +619,30 @@ section MainLoop
       st := { st with rpcSessions := st.rpcSessions.erase id }
 
   set st
+
+  -- Process the RPC-message and restart main-loop.
   match msg with
+  | Message.request id "shutdown" none =>
+    ctx.hOut.writeLspResponse ⟨id, Json.null⟩
+    mainLoop
   | Message.request id method (some params) =>
+    -- Requests are handled by the unmodified lean server.
     handleRequest id method (toJson params)
     mainLoop
   | Message.notification "exit" none =>
     let doc := st.doc
     doc.cancelTk.set
     return ()
-    | Message.request id "shutdown" none =>
-      ctx.hOut.writeLspResponse ⟨id, Json.null⟩
-      mainLoop
   | Message.notification method (some params) =>
+    -- Custom notification handler
     handleNotification method (toJson params)
     mainLoop
-    | _ => throwServerError s!"Got invalid JSON-RPC message: {toJson msg}"
+  | _ =>
+    throwServerError s!"Got invalid JSON-RPC message: {toJson msg}"
+
 end MainLoop
 
+
 def initAndRunWorker (i o e : FS.Stream) (opts : Options) (gameDir : String) : IO UInt32 := do
   let i ← maybeTee "fwIn.txt" false i
   let o ← maybeTee "fwOut.txt" true o
@@ -609,12 +681,13 @@ def initAndRunWorker (i o e : FS.Stream) (opts : Options) (gameDir : String) : I
     let levelInfo ← loadLevelData gameDir levelId.world levelId.level
     let some initializationOptions := initRequest.param.initializationOptions?
       | throwServerError "no initialization options found"
-    let gameWorkerState : GameWorkerState:= {
+    let gameWorkerState : WorkerState := {
       inventory := initializationOptions.inventory
       difficulty := initializationOptions.difficulty
       levelInfo
     }
     let (ctx, st) ← initializeWorker meta i o e initParams opts gameDir gameWorkerState
+    -- Run the main loop
     let _ ← StateRefT'.run (s := st) <| ReaderT.run (r := ctx) <|
       StateT.run (s := gameWorkerState) <| (mainLoop)
     return (0 : UInt32)
@@ -626,6 +699,11 @@ def initAndRunWorker (i o e : FS.Stream) (opts : Options) (gameDir : String) : I
       message := e.toString }] o
     return (1 : UInt32)
 
+/--
+The main function. Simply wrapping `initAndRunWorker`.
+
+TODO: The first arg `args[0]` is always expected to be `--server`. We could drop this completely.
+-/
 def workerMain (opts : Options) (args : List String): IO UInt32 := do
   let i ← IO.getStdin
   let o ← IO.getStdout
@@ -643,4 +721,4 @@ def workerMain (opts : Options) (args : List String): IO UInt32 := do
     e.putStrLn s!"worker initialization error: {err}"
     return (1 : UInt32)
 
-end MyServer.FileWorker
+end GameServer.FileWorker