feat: snapshot trees and language processors

CODEOWNERS

@@ -13,6 +13,7 @@
 /src/Lean/Data/Lsp/ @mhuisi
 /src/Lean/Elab/Deriving/ @semorrison
 /src/Lean/Elab/Tactic/ @semorrison
+/src/Lean/Language/ @Kha
 /src/Lean/Meta/Tactic/ @leodemoura
 /src/Lean/Parser/ @Kha
 /src/Lean/PrettyPrinter/ @Kha

src/Init/Data/Channel.lean

@@ -41,7 +41,7 @@ Sends a message on an `Channel`.
 
 This function does not block.
 -/
-def Channel.send (v : α) (ch : Channel α) : BaseIO Unit :=
+def Channel.send (ch : Channel α) (v : α) : BaseIO Unit :=
   ch.atomically do
     let st ← get
     if st.closed then return

src/Init/Data/Option/Basic.lean

@@ -37,6 +37,13 @@ def toMonad [Monad m] [Alternative m] : Option α → m α
   | none,   _ => none
   | some a, b => b a
 
+/-- Runs `f` on `o`'s value, if any, and returns its result, or else returns `none`.  -/
+@[inline] protected def bindM [Monad m] (f : α → m (Option β)) (o : Option α) : m (Option β) := do
+  if let some a := o then
+    return (← f a)
+  else
+    return none
+
 @[inline] protected def mapM [Monad m] (f : α → m β) (o : Option α) : m (Option β) := do
   if let some a := o then
     return some (← f a)

src/Lean/Data/JsonRpc.lean

@@ -183,6 +183,9 @@ structure ResponseError (α : Type u) where
 instance [ToJson α] : CoeOut (ResponseError α) Message :=
   ⟨fun r => Message.responseError r.id r.code r.message (r.data?.map toJson)⟩
 
+instance : CoeOut (ResponseError Unit) Message :=
+  ⟨fun r => Message.responseError r.id r.code r.message none⟩
+
 instance : Coe String RequestID := ⟨RequestID.str⟩
 instance : Coe JsonNumber RequestID := ⟨RequestID.num⟩
 

src/Lean/Data/Lsp/Internal.lean

@@ -84,6 +84,7 @@ structure RefInfo.Location where
   range       : Lsp.Range
   /-- Parent declaration of the reference. `none` if the reference is itself a declaration. -/
   parentDecl? : Option RefInfo.ParentDecl
+deriving Inhabited
 
 /-- Definition site and usage sites of a reference. Obtained from `Lean.Server.RefInfo`. -/
 structure RefInfo where

src/Lean/Data/Lsp/Ipc.lean

@@ -84,26 +84,32 @@ partial def readResponseAs (expectedID : RequestID) (α) [FromJson α] :
 def waitForExit : IpcM UInt32 := do
   (←read).wait
 
-/-- Waits for the worker to emit all diagnostics for the current document version
-and returns them as a list. -/
+/--
+Waits for the worker to emit all diagnostic notifications for the current document version and
+returns the last notification, if any.
+
+We used to return all notifications but with debouncing in the server, this would not be
+deterministic anymore as what messages are dropped depends on wall-clock timing.
+ -/
 partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target : DocumentUri) (version : Nat)
-: IpcM (List (Notification PublishDiagnosticsParams)) := do
+: IpcM (Option (Notification PublishDiagnosticsParams)) := do
   writeRequest ⟨waitForDiagnosticsId, "textDocument/waitForDiagnostics", WaitForDiagnosticsParams.mk target version⟩
-  let rec loop : IpcM (List (Notification PublishDiagnosticsParams)) := do
+  loop
+where
+  loop := do
     match (←readMessage) with
     | Message.response id _ =>
-      if id == waitForDiagnosticsId then return []
+      if id == waitForDiagnosticsId then return none
       else loop
     | Message.responseError id _    msg _ =>
       if id == waitForDiagnosticsId then
         throw $ userError s!"Waiting for diagnostics failed: {msg}"
       else loop
     | Message.notification "textDocument/publishDiagnostics" (some param) =>
       match fromJson? (toJson param) with
-      | Except.ok diagnosticParam => return ⟨"textDocument/publishDiagnostics", diagnosticParam⟩ :: (←loop)
+      | Except.ok diagnosticParam => return (← loop).getD ⟨"textDocument/publishDiagnostics", diagnosticParam⟩
       | Except.error inner => throw $ userError s!"Cannot decode publishDiagnostics parameters\n{inner}"
     | _ => loop
-  loop
 
 def runWith (lean : System.FilePath) (args : Array String := #[]) (test : IpcM α) : IO α := do
   let proc ← Process.spawn {

src/Lean/DeclarationRange.lean

@@ -48,14 +48,14 @@ def addDeclarationRanges [MonadEnv m] (declName : Name) (declRanges : Declaratio
 def findDeclarationRangesCore? [Monad m] [MonadEnv m] (declName : Name) : m (Option DeclarationRanges) :=
   return declRangeExt.find? (← getEnv) declName
 
-def findDeclarationRanges? [Monad m] [MonadEnv m] [MonadLiftT IO m] (declName : Name) : m (Option DeclarationRanges) := do
+def findDeclarationRanges? [Monad m] [MonadEnv m] [MonadLiftT BaseIO m] (declName : Name) : m (Option DeclarationRanges) := do
   let env ← getEnv
   let ranges ← if isAuxRecursor env declName || isNoConfusion env declName || (← isRec declName)  then
     findDeclarationRangesCore? declName.getPrefix
   else
     findDeclarationRangesCore? declName
   match ranges with
-  | none => return (← builtinDeclRanges.get (m := IO)).find? declName
+  | none => return (← builtinDeclRanges.get (m := BaseIO)).find? declName
   | some _ => return ranges
 
 end Lean

src/Lean/Elab/Command.lean

@@ -141,7 +141,8 @@ private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceStat
   let mut log := log
   let traces' := traces.toArray.qsort fun ((a, _), _) ((b, _), _) => a < b
   for ((pos, endPos), traceMsg) in traces' do
-    log := log.add <| mkMessageCore ctx.fileName ctx.fileMap (.joinSep traceMsg.toList "\n") .information pos endPos
+    let data := .tagged `_traceMsg <| .joinSep traceMsg.toList "\n"
+    log := log.add <| mkMessageCore ctx.fileName ctx.fileMap data .information pos endPos
   return log
 
 private def addTraceAsMessages : CommandElabM Unit := do
@@ -268,11 +269,6 @@ instance : MonadRecDepth CommandElabM where
   getRecDepth      := return (← read).currRecDepth
   getMaxRecDepth   := return (← get).maxRecDepth
 
-register_builtin_option showPartialSyntaxErrors : Bool := {
-  defValue := false
-  descr    := "show elaboration errors from partial syntax trees (i.e. after parser recovery)"
-}
-
 builtin_initialize registerTraceClass `Elab.command
 
 partial def elabCommand (stx : Syntax) : CommandElabM Unit := do
@@ -321,11 +317,6 @@ def elabCommandTopLevel (stx : Syntax) : CommandElabM Unit := withRef stx do pro
     -- note the order: first process current messages & info trees, then add back old messages & trees,
     -- then convert new traces to messages
     let mut msgs := (← get).messages
-    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check in general
-    if !showPartialSyntaxErrors.get (← getOptions) && initMsgs.hasErrors && stx.hasMissing then
-      -- discard elaboration errors, except for a few important and unlikely misleading ones, on parse error
-      msgs := ⟨msgs.msgs.filter fun msg =>
-        msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder || tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
     for tree in (← getInfoTrees) do
       trace[Elab.info] (← tree.format)
     modify fun st => { st with

src/Lean/Elab/Frontend.lean

@@ -4,8 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura, Sebastian Ullrich
 -/
 prelude
-import Lean.Elab.Import
-import Lean.Elab.Command
+import Lean.Language.Lean
 import Lean.Util.Profile
 import Lean.Server.References
 
@@ -40,7 +39,19 @@ def setCommandState (commandState : Command.State) : FrontendM Unit :=
 
 def elabCommandAtFrontend (stx : Syntax) : FrontendM Unit := do
   runCommandElabM do
+    let initMsgs ← modifyGet fun st => (st.messages, { st with messages := {} })
     Command.elabCommandTopLevel stx
+    let mut msgs := (← get).messages
+    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check
+    -- in general
+    if !Language.Lean.showPartialSyntaxErrors.get (← getOptions) && initMsgs.hasErrors &&
+        stx.hasMissing then
+      -- discard elaboration errors, except for a few important and unlikely misleading ones, on
+      -- parse error
+      msgs := ⟨msgs.msgs.filter fun msg =>
+        msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder ||
+          tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
+    modify ({ · with messages := initMsgs ++ msgs })
 
 def updateCmdPos : FrontendM Unit := do
   modify fun s => { s with cmdPos := s.parserState.pos }
@@ -86,12 +97,8 @@ def process (input : String) (env : Environment) (opts : Options) (fileName : Op
   pure (s.commandState.env, s.commandState.messages)
 
 builtin_initialize
-  registerOption `printMessageEndPos { defValue := false, descr := "print end position of each message in addition to start position" }
   registerTraceClass `Elab.info
 
-def getPrintMessageEndPos (opts : Options) : Bool :=
-  opts.getBool `printMessageEndPos false
-
 @[export lean_run_frontend]
 def runFrontend
     (input : String)
@@ -102,26 +109,50 @@ def runFrontend
     (ileanFileName? : Option String := none)
     : IO (Environment × Bool) := do
   let inputCtx := Parser.mkInputContext input fileName
-  let (header, parserState, messages) ← Parser.parseHeader inputCtx
-  -- allow `env` to be leaked, which would live until the end of the process anyway
-  let (env, messages) ← processHeader (leakEnv := true) header opts messages inputCtx trustLevel
-  let env := env.setMainModule mainModuleName
-  let mut commandState := Command.mkState env messages opts
-
-  if ileanFileName?.isSome then
-    -- Collect InfoTrees so we can later extract and export their info to the ilean file
-    commandState := { commandState with infoState.enabled := true }
-
-  let s ← IO.processCommands inputCtx parserState commandState
-  for msg in s.commandState.messages.toList do
-    IO.print (← msg.toString (includeEndPos := getPrintMessageEndPos opts))
-
+  -- TODO: replace with `#lang` processing
+  if /- Lean #lang? -/ true then
+    -- Temporarily keep alive old cmdline driver for the Lean language so that we don't pay the
+    -- overhead of passing the environment between snapshots until we actually make good use of it
+    -- outside the server
+    let (header, parserState, messages) ← Parser.parseHeader inputCtx
+    -- allow `env` to be leaked, which would live until the end of the process anyway
+    let (env, messages) ← processHeader (leakEnv := true) header opts messages inputCtx trustLevel
+    let env := env.setMainModule mainModuleName
+    let mut commandState := Command.mkState env messages opts
+
+    if ileanFileName?.isSome then
+      -- Collect InfoTrees so we can later extract and export their info to the ilean file
+      commandState := { commandState with infoState.enabled := true }
+
+    let s ← IO.processCommands inputCtx parserState commandState
+    for msg in s.commandState.messages.toList do
+      IO.print (← msg.toString (includeEndPos := Language.printMessageEndPos.get opts))
+
+    if let some ileanFileName := ileanFileName? then
+      let trees := s.commandState.infoState.trees.toArray
+      let references ←
+        Lean.Server.findModuleRefs inputCtx.fileMap trees (localVars := false) |>.toLspModuleRefs
+      let ilean := { module := mainModuleName, references : Lean.Server.Ilean }
+      IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
+
+    return (s.commandState.env, !s.commandState.messages.hasErrors)
+
+  let ctx := { inputCtx with mainModuleName, opts, trustLevel }
+  let processor := Language.Lean.process
+  let snap ← processor none ctx
+  let snaps := Language.toSnapshotTree snap
+  snaps.runAndReport opts
   if let some ileanFileName := ileanFileName? then
-    let trees := s.commandState.infoState.trees.toArray
+    let trees := snaps.getAll.concatMap (match ·.infoTree? with | some t => #[t] | _ => #[])
     let references := Lean.Server.findModuleRefs inputCtx.fileMap trees (localVars := false)
     let ilean := { module := mainModuleName, references := ← references.toLspModuleRefs : Lean.Server.Ilean }
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
-  pure (s.commandState.env, !s.commandState.messages.hasErrors)
+  let hasErrors := snaps.getAll.any (·.diagnostics.msgLog.hasErrors)
+  -- TODO: remove default when reworking cmdline interface in Lean; currently the only case
+  -- where we use the environment despite errors in the file is `--stats`
+  let env := Language.Lean.waitForFinalEnv? snap |>.getD (← mkEmptyEnvironment)
+  pure (env, !hasErrors)
+
 
 end Lean.Elab

src/Lean/Elab/InfoTree/Main.lean

@@ -94,16 +94,18 @@ partial def InfoTree.substitute (tree : InfoTree) (assignment : PersistentHashMa
     | none      => hole id
     | some tree => substitute tree assignment
 
-def ContextInfo.runMetaM (info : ContextInfo) (lctx : LocalContext) (x : MetaM α) : IO α := do
-  let x := x.run { lctx := lctx } { mctx := info.mctx }
+/-- Embeds a `CoreM` action in `IO` by supplying the information stored in `info`. -/
+def ContextInfo.runCoreM (info : ContextInfo) (x : CoreM α) : IO α := do
   /-
     We must execute `x` using the `ngen` stored in `info`. Otherwise, we may create `MVarId`s and `FVarId`s that
     have been used in `lctx` and `info.mctx`.
   -/
-  let ((a, _), _) ←
+  (·.1) <$>
     x.toIO { options := info.options, currNamespace := info.currNamespace, openDecls := info.openDecls, fileName := "<InfoTree>", fileMap := default }
            { env := info.env, ngen := info.ngen }
-  return a
+
+def ContextInfo.runMetaM (info : ContextInfo) (lctx : LocalContext) (x : MetaM α) : IO α := do
+  (·.1) <$> info.runCoreM (x.run { lctx := lctx } { mctx := info.mctx })
 
 def ContextInfo.toPPContext (info : ContextInfo) (lctx : LocalContext) : PPContext :=
   { env  := info.env, mctx := info.mctx, lctx := lctx,