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,104 +59,146 @@ 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.
-  0: no check
-  1: give warnings
-  2: give errors
+  Difficulty determines whether tactics/theorems can be locked.
+  * 0: do not check
+  * 1: give warnings when locked items are used
+  * 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
+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
-      messages := st.messages.add {
-        fileName := inputCtx.fileName
-        severity := MessageSeverity.error
-        pos      := inputCtx.fileMap.toPosition (info.getPos?.getD 0)
-        data     := s
-      }
-    }
+    messages := st.messages.add {
+      fileName := inputCtx.fileName
+      severity := MessageSeverity.error
+      pos      := inputCtx.fileMap.toPosition (info.getPos?.getD 0)
+      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. -/
-partial def findForbiddenTactics (inputCtx : Parser.InputContext)
-    (gameWorkerState : GameWorkerState) (stx : Syntax) :
-    Elab.Command.CommandElabM Unit := do
-  let levelInfo := gameWorkerState.levelInfo
+/--
+Find all tactics in syntax object that are forbidden according to a
+set `allowed` of allowed tactics.
+-/
+partial def findForbiddenTactics (inputCtx : Parser.InputContext) (workerState : WorkerState)
+    (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.
-    -- Note: We need a whitelist because we cannot really distinguish keywords from tactics.
+    -- Atoms might be tactic names or other keywords.
+    -- 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.
-        match gameWorkerState.inventory.find? (· == val) with
+        -- Tactic will never be introduced in the game.
+        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!"
-        | some _ => pure () -- tactic is in the inventory, allow it.
+          -- Tactic is not in the inventory.
+          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!"
-          | some _ => pure () -- tactic is in the inventory, allow it.
+            -- Tactic is marked as locked and not in the inventory.
+            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
-        -- Forbid the theorem we are proving currently
+          -- Not a theorem, no checks needed.
+          | return ()
         if some n = levelInfo.statementName then
-          addErrorMessage info inputCtx s!"Structural recursion: you can't use '{n}' to proof itself!"
-
-        let lemmasAndDefs := levelInfo.lemmas ++ levelInfo.definitions
-        match lemmasAndDefs.find? (fun l => l.name == n) with
-        | none => addWarningMessage info s!"You have not unlocked the lemma/definition '{n}' yet!"
-        | some lem =>
-          if lem.locked then
-            addWarningMessage info s!"You have not unlocked the lemma/definition '{n}' yet!"
-          else if lem.disabled then
-            addWarningMessage info s!"The lemma/definition '{n}' is disabled in this level!"
-where addWarningMessage (info : SourceInfo) (s : MessageData) :=
-  let difficulty := gameWorkerState.difficulty
+          -- Forbid the theorem we are proving currently
+          addMessage info inputCtx (severity := .error)
+            s!"Structural recursion: you can't use '{n}' to proof itself!"
+        let theoremsAndDefs := levelInfo.lemmas ++ levelInfo.definitions
+        match theoremsAndDefs.find? (·.name == n) with
+        | none =>
+          -- Theorem will never be introduced in this game
+          addMessageByDifficulty info s!"You have not unlocked the theorem/definition '{n}' yet!"
+        | some thm =>
+          -- Theorem is introduced at some point in the game.
+          if thm.locked then
+            -- 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
-      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
-      }
-    }
+    addMessage info inputCtx (if difficulty > 1 then .error else .warning) 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,226 +395,254 @@ end Elab
 section Updates
 
 /-- Given the new document, updates editable doc state. -/
-  def updateDocument (newMeta : DocumentMeta) : GameWorkerM Unit := do
-    let s ← get
-    let ctx ← read
-    let oldDoc := (← StateT.lift get).doc
-    oldDoc.cancelTk.set
-    let initHeaderStx := (← StateT.lift get).initHeaderStx
-    let (newHeaderStx, newMpState, _) ← Parser.parseHeader newMeta.mkInputContext
-    let cancelTk ← CancelToken.new
-    let headSnapTask := oldDoc.cmdSnaps.waitHead?
-    let newSnaps ← if initHeaderStx != newHeaderStx then
-      EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
-        IO.sleep ctx.initParams.editDelay.toUInt32
-        cancelTk.check
-        IO.Process.exit 2
-    else EIO.mapTask (ε := ElabTaskError) (t := headSnapTask) (prio := .dedicated) fun headSnap?? => do
-      -- There is always at least one snapshot absent exceptions
-      let some headSnap ← MonadExcept.ofExcept headSnap?? | panic! "empty snapshots"
-      let newHeaderSnap := { headSnap with stx := newHeaderStx, mpState := newMpState }
-      let changePos := oldDoc.meta.text.source.firstDiffPos newMeta.text.source
-      -- Ignore exceptions, we are only interested in the successful snapshots
-      let (cmdSnaps, _) ← oldDoc.cmdSnaps.getFinishedPrefix
-      -- 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 ← pure (cmdSnaps.takeWhile (fun s => s.endPos < changePos))
-      if h : validSnaps.length ≤ 1 then
-        validSnaps := [newHeaderSnap]
-      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. -/
-        have : validSnaps.length ≥ 2 := Nat.gt_of_not_le h
-        let mut lastSnap := validSnaps.getLast (by subst ·; simp at h)
-        let preLastSnap :=
-          have : 0 < validSnaps.length := Nat.lt_of_lt_of_le (by decide) this
-          have : validSnaps.length - 2 < validSnaps.length := Nat.sub_lt this (by decide)
-          validSnaps[validSnaps.length - 2]
-        let newLastStx ← parseNextCmd newMeta.mkInputContext preLastSnap
-        if newLastStx != lastSnap.stx then
-          validSnaps := validSnaps.dropLast
-      -- wait for a bit, giving the initial `cancelTk.check` in `nextCmdSnap` time to trigger
-      -- before kicking off any expensive elaboration (TODO: make expensive elaboration cancelable)
-      unfoldSnaps newMeta validSnaps.toArray cancelTk s ctx
-        (startAfterMs := ctx.initParams.editDelay.toUInt32)
-    StateT.lift <| modify fun st => { st with
-      doc := { meta := newMeta, cmdSnaps := AsyncList.delayed newSnaps, cancelTk }}
+def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := do
+  let s ← get
+  let ctx ← read
+  let oldDoc := (← StateT.lift get).doc
+  oldDoc.cancelTk.set
+  let initHeaderStx := (← StateT.lift get).initHeaderStx
+  let (newHeaderStx, newMpState, _) ← Parser.parseHeader newMeta.mkInputContext
+  let cancelTk ← CancelToken.new
+  let headSnapTask := oldDoc.cmdSnaps.waitHead?
+  let newSnaps ← if initHeaderStx != newHeaderStx then
+    EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
+      IO.sleep ctx.initParams.editDelay.toUInt32
+      cancelTk.check
+      IO.Process.exit 2
+  else EIO.mapTask (ε := ElabTaskError) (t := headSnapTask) (prio := .dedicated) fun headSnap?? => do
+    -- There is always at least one snapshot absent exceptions
+    let some headSnap ← MonadExcept.ofExcept headSnap?? | panic! "empty snapshots"
+    let newHeaderSnap := { headSnap with stx := newHeaderStx, mpState := newMpState }
+    let changePos := oldDoc.meta.text.source.firstDiffPos newMeta.text.source
+    -- Ignore exceptions, we are only interested in the successful snapshots
+    let (cmdSnaps, _) ← oldDoc.cmdSnaps.getFinishedPrefix
+    -- 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 ← pure (cmdSnaps.takeWhile (fun s => s.endPos < changePos))
+    if h : validSnaps.length ≤ 1 then
+      validSnaps := [newHeaderSnap]
+    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. -/
+      have : validSnaps.length ≥ 2 := Nat.gt_of_not_le h
+      let mut lastSnap := validSnaps.getLast (by subst ·; simp at h)
+      let preLastSnap :=
+        have : 0 < validSnaps.length := Nat.lt_of_lt_of_le (by decide) this
+        have : validSnaps.length - 2 < validSnaps.length := Nat.sub_lt this (by decide)
+        validSnaps[validSnaps.length - 2]
+      let newLastStx ← parseNextCmd newMeta.mkInputContext preLastSnap
+      if newLastStx != lastSnap.stx then
+        validSnaps := validSnaps.dropLast
+    -- wait for a bit, giving the initial `cancelTk.check` in `nextCmdSnap` time to trigger
+    -- before kicking off any expensive elaboration (TODO: make expensive elaboration cancelable)
+    unfoldSnaps newMeta validSnaps.toArray cancelTk s ctx
+      (startAfterMs := ctx.initParams.editDelay.toUInt32)
+  StateT.lift <| modify fun st => { st with
+    doc := { meta := newMeta, cmdSnaps := AsyncList.delayed newSnaps, cancelTk }}
 
 end Updates
 
 section Initialization
 
-  def DocumentMeta.mkInputContext (doc : DocumentMeta) : Parser.InputContext where
-    input    := "" -- No header!
-    fileName := (System.Uri.fileUriToPath? doc.uri).getD doc.uri |>.toString
-    fileMap  := default
-
-  def compileHeader (m : DocumentMeta) (hOut : FS.Stream) (opts : Options) (hasWidgets : Bool)
-      (gameDir : String) (module : Name):
-      IO (Syntax × Task (Except Error (Snapshot × SearchPath))) := do
-    -- Determine search paths of the game project by running `lake env printenv LEAN_PATH`.
-    let out ← IO.Process.output
-      { cwd := gameDir, cmd := "lake", args := #["env","printenv","LEAN_PATH"] }
-    if out.exitCode != 0 then
-      throwServerError s!"Error while running Lake: {out.stderr}"
-
-    -- Make the paths relative to the current directory
-    let paths : List System.FilePath := System.SearchPath.parse out.stdout.trim
-    let currentDir ← IO.currentDir
-    let paths := paths.map fun p => currentDir / (gameDir : System.FilePath) / p
-
-    -- Set the search path
-    Lean.searchPathRef.set paths
-
-    let env ← importModules' #[{ module := `Init : Import }, { module := module : Import }]
-
-    -- use empty header
-    let (headerStx, headerParserState, msgLog) ← Parser.parseHeader
-      {m.mkInputContext with
-        input := ""
-        fileMap := FileMap.ofString ""}
-    (headerStx, ·) <$> EIO.asTask do
-      let mut srcSearchPath : SearchPath := paths --← initSrcSearchPath (← getBuildDir)
-      let headerEnv := env
-      let mut headerEnv := headerEnv
-      try
-        if let some path := System.Uri.fileUriToPath? m.uri then
-          headerEnv := headerEnv.setMainModule (← moduleNameOfFileName path none)
-      catch _ => pure ()
-      let cmdState := Elab.Command.mkState headerEnv {} opts
-      let cmdState := { cmdState with infoState := {
-        enabled := true
-        trees := #[Elab.InfoTree.context ({
-          env     := headerEnv
-          fileMap := m.text
-          ngen    := { namePrefix := `_worker }
-        }) (Elab.InfoTree.node
-            (Elab.Info.ofCommandInfo { elaborator := `header, stx := headerStx })
-            (headerStx[1].getArgs.toList.map (fun importStx =>
-              Elab.InfoTree.node (Elab.Info.ofCommandInfo {
-                elaborator := `import
-                stx := importStx
-              }) #[].toPArray'
-            )).toPArray'
-        )].toPArray'
-      }}
-      let headerSnap := {
-        beginPos := 0
-        stx := headerStx
-        mpState := {}
-        cmdState := cmdState
-        interactiveDiags := ← cmdState.messages.msgs.mapM (Widget.msgToInteractiveDiagnostic m.text · hasWidgets)
-        tacticCache := (← IO.mkRef {})
-      }
-      publishDiagnostics m headerSnap.diagnostics.toArray hOut
-      return (headerSnap, srcSearchPath)
-
-  def initializeWorker (meta : DocumentMeta) (i o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
-      (gameDir : String) (gameWorkerState : GameWorkerState) : IO (WorkerContext × WorkerState) := do
-    let clientHasWidgets := initParams.initializationOptions?.bind (·.hasWidgets?) |>.getD false
-
-    let (headerStx, headerTask) ← compileHeader meta o opts (hasWidgets := clientHasWidgets)
-      gameDir gameWorkerState.levelInfo.module
-    let cancelTk ← CancelToken.new
-    let ctx :=
-      { hIn  := i
-        hOut := o
-        hLog := e
-        headerTask
-        initParams
-        clientHasWidgets
-      }
-    let cmdSnaps ← EIO.mapTask (t := headerTask) (match · with
-      | Except.ok (s, _) => unfoldSnaps meta #[s] cancelTk gameWorkerState ctx (startAfterMs := 0)
-      | Except.error e   => throw (e : ElabTaskError))
-    let doc : EditableDocument := { meta, cmdSnaps := AsyncList.delayed cmdSnaps, cancelTk }
-    return (ctx,
-    { doc             := doc
-      initHeaderStx   := headerStx
-      currHeaderStx      := headerStx
-      importCachingTask? := none
-      pendingRequests := RBMap.empty
-      rpcSessions     := RBMap.empty
-    })
+def DocumentMeta.mkInputContext (doc : DocumentMeta) : Parser.InputContext where
+  input    := "" -- No header!
+  fileName := (System.Uri.fileUriToPath? doc.uri).getD doc.uri |>.toString
+  fileMap  := default
+
+def compileHeader (m : DocumentMeta) (hOut : FS.Stream) (opts : Options) (hasWidgets : Bool)
+    (gameDir : String) (module : Name):
+    IO (Syntax × Task (Except Error (Snapshot × SearchPath))) := do
+  -- Determine search paths of the game project by running `lake env printenv LEAN_PATH`.
+  let out ← IO.Process.output
+    { cwd := gameDir, cmd := "lake", args := #["env","printenv","LEAN_PATH"] }
+  if out.exitCode != 0 then
+    throwServerError s!"Error while running Lake: {out.stderr}"
+
+  -- Make the paths relative to the current directory
+  let paths : List System.FilePath := System.SearchPath.parse out.stdout.trim
+  let currentDir ← IO.currentDir
+  let paths := paths.map fun p => currentDir / (gameDir : System.FilePath) / p
+
+  -- Set the search path
+  Lean.searchPathRef.set paths
+
+  let env ← importModules' #[{ module := `Init : Import }, { module := module : Import }]
+
+  -- use empty header
+  let (headerStx, headerParserState, msgLog) ← Parser.parseHeader
+    {m.mkInputContext with
+      input := ""
+      fileMap := FileMap.ofString ""}
+  (headerStx, ·) <$> EIO.asTask do
+    let mut srcSearchPath : SearchPath := paths --← initSrcSearchPath (← getBuildDir)
+    let headerEnv := env
+    let mut headerEnv := headerEnv
+    try
+      if let some path := System.Uri.fileUriToPath? m.uri then
+        headerEnv := headerEnv.setMainModule (← moduleNameOfFileName path none)
+    catch _ => pure ()
+    let cmdState := Elab.Command.mkState headerEnv {} opts
+    let cmdState := { cmdState with infoState := {
+      enabled := true
+      trees := #[Elab.InfoTree.context ({
+        env     := headerEnv
+        fileMap := m.text
+        ngen    := { namePrefix := `_worker }
+      }) (Elab.InfoTree.node
+          (Elab.Info.ofCommandInfo { elaborator := `header, stx := headerStx })
+          (headerStx[1].getArgs.toList.map (fun importStx =>
+            Elab.InfoTree.node (Elab.Info.ofCommandInfo {
+              elaborator := `import
+              stx := importStx
+            }) #[].toPArray'
+          )).toPArray'
+      )].toPArray'
+    }}
+    let headerSnap := {
+      beginPos := 0
+      stx := headerStx
+      mpState := {}
+      cmdState := cmdState
+      interactiveDiags := ← cmdState.messages.msgs.mapM (Widget.msgToInteractiveDiagnostic m.text · hasWidgets)
+      tacticCache := (← IO.mkRef {})
+    }
+    publishDiagnostics m headerSnap.diagnostics.toArray hOut
+    return (headerSnap, srcSearchPath)
+
+def initializeWorker (meta : DocumentMeta) (i o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
+    (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)
+    gameDir gameWorkerState.levelInfo.module
+  let cancelTk ← CancelToken.new
+  let ctx :=
+    { hIn  := i
+      hOut := o
+      hLog := e
+      headerTask
+      initParams
+      clientHasWidgets
+    }
+  let cmdSnaps ← EIO.mapTask (t := headerTask) (match · with
+    | Except.ok (s, _) => unfoldSnaps meta #[s] cancelTk gameWorkerState ctx (startAfterMs := 0)
+    | Except.error e   => throw (e : ElabTaskError))
+  let doc : EditableDocument := { meta, cmdSnaps := AsyncList.delayed cmdSnaps, cancelTk }
+  return (ctx,
+  { doc             := doc
+    initHeaderStx   := headerStx
+    currHeaderStx      := headerStx
+    importCachingTask? := none
+    pendingRequests := RBMap.empty
+    rpcSessions     := RBMap.empty
+  })
 
 end Initialization
 
 section NotificationHandling
 
-  def handleDidChange (p : DidChangeTextDocumentParams) : GameWorkerM Unit := do
-    let docId := p.textDocument
-    let changes := p.contentChanges
-    let oldDoc := (← StateT.lift get).doc
-    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 := foldDocumentChanges changes oldDoc.meta.text
-      updateDocument ⟨docId.uri, newVersion, newDocText, .always⟩
+def handleDidChange (p : DidChangeTextDocumentParams) : WorkerM Unit := do
+  let docId := p.textDocument
+  let changes := p.contentChanges
+  let oldDoc := (← StateT.lift get).doc
+  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 := foldDocumentChanges changes oldDoc.meta.text
+    updateDocument ⟨docId.uri, newVersion, newDocText, .always⟩
 
 end NotificationHandling
 
 section MessageHandling
 
-  def handleNotification (method : String) (params : Json) : GameWorkerM Unit := do
-    let handle := fun paramType [FromJson paramType] (handler : paramType → GameWorkerM Unit) =>
-      (StateT.lift <| parseParams paramType params) >>= handler
-    match method with
-    | "textDocument/didChange" => handle DidChangeTextDocumentParams (handleDidChange)
-    | "$/cancelRequest"        => handle CancelParams (handleCancelRequest ·)
-    | "$/setTrace"        => pure ()
-    | "$/lean/rpc/release"     => handle RpcReleaseParams (handleRpcRelease ·)
-    | "$/lean/rpc/keepAlive"   => handle RpcKeepAliveParams (handleRpcKeepAlive ·)
-    | _                        => throwServerError s!"Got unsupported notification method: {method}"
+
+/--
+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 ·)
+  -- 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
-    let msg ← ctx.hIn.readLspMessage
-    let filterFinishedTasks (acc : PendingRequestMap) (id : RequestID) (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. -/
-        if let Except.error e := task.get then
-          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 }
-
-    -- 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
-      if (← sesh.hasExpired) then
-        st := { st with rpcSessions := st.rpcSessions.erase id }
-
-    set st
-    match msg with
-    | Message.request id method (some params) =>
-      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) =>
-      handleNotification method (toJson params)
-      mainLoop
-    | _ => throwServerError s!"Got invalid JSON-RPC message: {toJson msg}"
+
+/--
+Erase finished tasks if there are no errors.
+-/
+private def filterFinishedTasks (acc : PendingRequestMap) (id : RequestID)
+    (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. -/
+    if let Except.error e := task.get then
+      throwServerError s!"Failed responding to request {id}: {e}"
+    pure <| acc.erase id
+  else pure acc
+
+/--
+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
+    if (← sesh.hasExpired) then
+      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.notification method (some params) =>
+    -- Custom notification handler
+    handleNotification method (toJson params)
+    mainLoop
+  | _ =>
+    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