feat: creation and reporting for asynchronous elaboration tasks (#6170)

This PR adds core metaprogramming functions for forking off background
tasks from elaboration such that their results are visible to reporting
and the language server

src/Lean/CoreM.lean

@@ -11,6 +11,7 @@ import Lean.ResolveName
 import Lean.Elab.InfoTree.Types
 import Lean.MonadEnv
 import Lean.Elab.Exception
+import Lean.Language.Basic
 
 namespace Lean
 register_builtin_option diagnostics : Bool := {
@@ -72,6 +73,13 @@ structure State where
   messages        : MessageLog     := {}
   /-- Info tree. We have the info tree here because we want to update it while adding attributes. -/
   infoState       : Elab.InfoState := {}
+  /--
+  Snapshot trees of asynchronous subtasks. As these are untyped and reported only at the end of the
+  command's main elaboration thread, they are only useful for basic message log reporting; for
+  incremental reporting and reuse within a long-running elaboration thread, types rooted in
+  `CommandParsedSnapshot` need to be adjusted.
+  -/
+  snapshotTasks : Array (Language.SnapshotTask Language.SnapshotTree) := #[]
   deriving Nonempty
 
 /-- Context for the CoreM monad. -/
@@ -180,7 +188,8 @@ instance : Elab.MonadInfoTree CoreM where
   modifyInfoState f := modify fun s => { s with infoState := f s.infoState }
 
 @[inline] def modifyCache (f : Cache → Cache) : CoreM Unit :=
-  modify fun ⟨env, next, ngen, trace, cache, messages, infoState⟩ => ⟨env, next, ngen, trace, f cache, messages, infoState⟩
+  modify fun ⟨env, next, ngen, trace, cache, messages, infoState, snaps⟩ =>
+   ⟨env, next, ngen, trace, f cache, messages, infoState, snaps⟩
 
 @[inline] def modifyInstLevelTypeCache (f : InstantiateLevelCache → InstantiateLevelCache) : CoreM Unit :=
   modifyCache fun ⟨c₁, c₂⟩ => ⟨f c₁, c₂⟩
@@ -355,13 +364,83 @@ instance : MonadLog CoreM where
     if (← read).suppressElabErrors then
       -- discard elaboration errors, except for a few important and unlikely misleading ones, on
       -- parse error
-      unless msg.data.hasTag (· matches `Elab.synthPlaceholder | `Tactic.unsolvedGoals) do
+      unless msg.data.hasTag (· matches `Elab.synthPlaceholder | `Tactic.unsolvedGoals | `trace) do
         return
 
     let ctx ← read
     let msg := { msg with data := MessageData.withNamingContext { currNamespace := ctx.currNamespace, openDecls := ctx.openDecls } msg.data };
     modify fun s => { s with messages := s.messages.add msg }
 
+/--
+Includes a given task (such as from `wrapAsyncAsSnapshot`) in the overall snapshot tree for this
+command's elaboration, making its result available to reporting and the language server. The
+reporter will not know about this snapshot tree node until the main elaboration thread for this
+command has finished so this function is not useful for incremental reporting within a longer
+elaboration thread but only for tasks that outlive it such as background kernel checking or proof
+elaboration.
+-/
+def logSnapshotTask (task : Language.SnapshotTask Language.SnapshotTree) : CoreM Unit :=
+  modify fun s => { s with snapshotTasks := s.snapshotTasks.push task }
+
+/-- Wraps the given action for use in `EIO.asTask` etc., discarding its final monadic state. -/
+def wrapAsync (act : Unit → CoreM α) : CoreM (EIO Exception α) := do
+  let st ← get
+  let ctx ← read
+  let heartbeats := (← IO.getNumHeartbeats) - ctx.initHeartbeats
+  return withCurrHeartbeats (do
+      -- include heartbeats since start of elaboration in new thread as well such that forking off
+      -- an action doesn't suddenly allow it to succeed from a lower heartbeat count
+      IO.addHeartbeats heartbeats.toUInt64
+      act () : CoreM _)
+    |>.run' ctx st
+
+/-- Option for capturing output to stderr during elaboration. -/
+register_builtin_option stderrAsMessages : Bool := {
+  defValue := true
+  group    := "server"
+  descr    := "(server) capture output to the Lean stderr channel (such as from `dbg_trace`) during elaboration of a command as a diagnostic message"
+}
+
+open Language in
+/--
+Wraps the given action for use in `BaseIO.asTask` etc., discarding its final state except for
+`logSnapshotTask` tasks, which are reported as part of the returned tree.
+-/
+def wrapAsyncAsSnapshot (act : Unit → CoreM Unit) (desc : String := by exact decl_name%.toString) :
+    CoreM (BaseIO SnapshotTree) := do
+  let t ← wrapAsync fun _ => do
+    IO.FS.withIsolatedStreams (isolateStderr := stderrAsMessages.get (← getOptions)) do
+      let tid ← IO.getTID
+      -- reset trace state and message log so as not to report them twice
+      modify ({ · with messages := {}, traceState := { tid } })
+      try
+        withTraceNode `Elab.async (fun _ => return desc) do
+          act ()
+      catch e =>
+        logError e.toMessageData
+      finally
+        addTraceAsMessages
+      get
+  let ctx ← readThe Core.Context
+  return do
+    match (← t.toBaseIO) with
+    | .ok (output, st) =>
+      let mut msgs := st.messages
+      if !output.isEmpty then
+        msgs := msgs.add {
+          fileName := ctx.fileName
+          severity := MessageSeverity.information
+          pos      := ctx.fileMap.toPosition <| ctx.ref.getPos?.getD 0
+          data     := output
+        }
+      return .mk {
+        desc
+        diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog msgs)
+        traces := st.traceState
+      } st.snapshotTasks
+    -- interrupt or abort exception as `try catch` above should have caught any others
+    | .error _ => default
+
 end Core
 
 export Core (CoreM mkFreshUserName checkSystem withCurrHeartbeats)

src/Lean/Elab/Command.lean

@@ -84,6 +84,7 @@ structure State where
   ngen           : NameGenerator := {}
   infoState      : InfoState := {}
   traceState     : TraceState := {}
+  snapshotTasks  : Array (Language.SnapshotTask Language.SnapshotTree) := #[]
   deriving Nonempty
 
 structure Context where
@@ -114,8 +115,7 @@ structure Context where
   -/
   suppressElabErrors : Bool := false
 
-abbrev CommandElabCoreM (ε) := ReaderT Context $ StateRefT State $ EIO ε
-abbrev CommandElabM := CommandElabCoreM Exception
+abbrev CommandElabM := ReaderT Context $ StateRefT State $ EIO Exception
 abbrev CommandElab  := Syntax → CommandElabM Unit
 structure Linter where
   run : Syntax → CommandElabM Unit
@@ -198,36 +198,6 @@ instance : AddErrorMessageContext CommandElabM where
     let msg ← addMacroStack msg ctx.macroStack
     return (ref, msg)
 
-def mkMessageAux (ctx : Context) (ref : Syntax) (msgData : MessageData) (severity : MessageSeverity) : Message :=
-  let pos := ref.getPos?.getD ctx.cmdPos
-  let endPos := ref.getTailPos?.getD pos
-  mkMessageCore ctx.fileName ctx.fileMap msgData severity pos endPos
-
-private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceState : TraceState) : MessageLog := Id.run do
-  if traceState.traces.isEmpty then return log
-  let mut traces : Std.HashMap (String.Pos × String.Pos) (Array MessageData) := ∅
-  for traceElem in traceState.traces do
-    let ref := replaceRef traceElem.ref ctx.ref
-    let pos := ref.getPos?.getD 0
-    let endPos := ref.getTailPos?.getD pos
-    traces := traces.insert (pos, endPos) <| traces.getD (pos, endPos) #[] |>.push traceElem.msg
-  let mut log := log
-  let traces' := traces.toArray.qsort fun ((a, _), _) ((b, _), _) => a < b
-  for ((pos, endPos), traceMsg) in traces' do
-    let data := .tagged `trace <| .joinSep traceMsg.toList "\n"
-    log := log.add <| mkMessageCore ctx.fileName ctx.fileMap data .information pos endPos
-  return log
-
-private def addTraceAsMessages : CommandElabM Unit := do
-  let ctx ← read
-  -- do not add trace messages if `trace.profiler.output` is set as it would be redundant and
-  -- pretty printing the trace messages is expensive
-  if trace.profiler.output.get? (← getOptions) |>.isNone then
-    modify fun s => { s with
-      messages          := addTraceAsMessagesCore ctx s.messages s.traceState
-      traceState.traces := {}
-    }
-
 private def runCore (x : CoreM α) : CommandElabM α := do
   let s ← get
   let ctx ← read
@@ -253,6 +223,7 @@ private def runCore (x : CoreM α) : CommandElabM α := do
     nextMacroScope := s.nextMacroScope
     infoState.enabled := s.infoState.enabled
     traceState := s.traceState
+    snapshotTasks := s.snapshotTasks
   }
   let (ea, coreS) ← liftM x
   modify fun s => { s with
@@ -261,17 +232,14 @@ private def runCore (x : CoreM α) : CommandElabM α := do
     ngen              := coreS.ngen
     infoState.trees   := s.infoState.trees.append coreS.infoState.trees
     traceState.traces := coreS.traceState.traces.map fun t => { t with ref := replaceRef t.ref ctx.ref }
+    snapshotTasks     := coreS.snapshotTasks
     messages          := s.messages ++ coreS.messages
   }
   return ea
 
 def liftCoreM (x : CoreM α) : CommandElabM α := do
   MonadExcept.ofExcept (← runCore (observing x))
 
-private def ioErrorToMessage (ctx : Context) (ref : Syntax) (err : IO.Error) : Message :=
-  let ref := getBetterRef ref ctx.macroStack
-  mkMessageAux ctx ref (toString err) MessageSeverity.error
-
 @[inline] def liftIO {α} (x : IO α) : CommandElabM α := do
   let ctx ← read
   IO.toEIO (fun (ex : IO.Error) => Exception.error ctx.ref ex.toString) x
@@ -294,9 +262,8 @@ instance : MonadLog CommandElabM where
   logMessage msg := do
     if (← read).suppressElabErrors then
       -- discard elaboration errors on parse error
-      -- NOTE: unlike `CoreM`'s `logMessage`, we do not currently have any command-level errors that
-      -- we want to allowlist
-      return
+      unless msg.data.hasTag (· matches `trace) do
+        return
     let currNamespace ← getCurrNamespace
     let openDecls ← getOpenDecls
     let msg := { msg with data := MessageData.withNamingContext { currNamespace := currNamespace, openDecls := openDecls } msg.data }
@@ -322,6 +289,61 @@ def runLinters (stx : Syntax) : CommandElabM Unit := do
             finally
               modify fun s => { savedState with messages := s.messages }
 
+/--
+Catches and logs exceptions occurring in `x`. Unlike `try catch` in `CommandElabM`, this function
+catches interrupt exceptions as well and thus is intended for use at the top level of elaboration.
+Interrupt and abort exceptions are caught but not logged.
+-/
+@[inline] def withLoggingExceptions (x : CommandElabM Unit) : CommandElabM Unit := fun ctx ref =>
+  EIO.catchExceptions (withLogging x ctx ref) (fun _ => pure ())
+
+@[inherit_doc Core.wrapAsync]
+def wrapAsync (act : Unit → CommandElabM α) : CommandElabM (EIO Exception α) := do
+  return act () |>.run (← read) |>.run' (← get)
+
+open Language in
+@[inherit_doc Core.wrapAsyncAsSnapshot]
+-- `CoreM` and `CommandElabM` are too different to meaningfully share this code
+def wrapAsyncAsSnapshot (act : Unit → CommandElabM Unit)
+    (desc : String := by exact decl_name%.toString) :
+    CommandElabM (BaseIO SnapshotTree) := do
+  let t ← wrapAsync fun _ => do
+    IO.FS.withIsolatedStreams (isolateStderr := Core.stderrAsMessages.get (← getOptions)) do
+      let tid ← IO.getTID
+      -- reset trace state and message log so as not to report them twice
+      modify ({ · with messages := {}, traceState := { tid } })
+      try
+        withTraceNode `Elab.async (fun _ => return desc) do
+          act ()
+      catch e =>
+        logError e.toMessageData
+      finally
+        addTraceAsMessages
+      get
+  let ctx ← read
+  return do
+    match (← t.toBaseIO) with
+    | .ok (output, st) =>
+      let mut msgs := st.messages
+      if !output.isEmpty then
+        msgs := msgs.add {
+          fileName := ctx.fileName
+          severity := MessageSeverity.information
+          pos      := ctx.fileMap.toPosition <| ctx.ref.getPos?.getD 0
+          data     := output
+        }
+      return .mk {
+        desc
+        diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog msgs)
+        traces := st.traceState
+      } st.snapshotTasks
+    -- interrupt or abort exception as `try catch` above should have caught any others
+    | .error _ => default
+
+@[inherit_doc Core.logSnapshotTask]
+def logSnapshotTask (task : Language.SnapshotTask Language.SnapshotTree) : CommandElabM Unit :=
+  modify fun s => { s with snapshotTasks := s.snapshotTasks.push task }
+
 protected def getCurrMacroScope : CommandElabM Nat  := do pure (← read).currMacroScope
 protected def getMainModule     : CommandElabM Name := do pure (← getEnv).mainModule
 
@@ -532,12 +554,6 @@ def elabCommandTopLevel (stx : Syntax) : CommandElabM Unit := withRef stx do pro
     let mut msgs := (← get).messages
     for tree in (← getInfoTrees) do
       trace[Elab.info] (← tree.format)
-    if (← isTracingEnabledFor `Elab.snapshotTree) then
-      if let some snap := (← read).snap? then
-        -- We can assume that the root command snapshot is not involved in parallelism yet, so this
-        -- should be true iff the command supports incrementality
-        if (← IO.hasFinished snap.new.result) then
-          liftCoreM <| Language.ToSnapshotTree.toSnapshotTree snap.new.result.get |>.trace
     modify fun st => { st with
       messages := initMsgs ++ msgs
       infoState := { st.infoState with trees := initInfoTrees ++ st.infoState.trees }
@@ -668,14 +684,6 @@ def runTermElabM (elabFn : Array Expr → TermElabM α) : CommandElabM α := do
           Term.addAutoBoundImplicits' xs someType fun xs _ =>
             Term.withoutAutoBoundImplicit <| elabFn xs
 
-/--
-Catches and logs exceptions occurring in `x`. Unlike `try catch` in `CommandElabM`, this function
-catches interrupt exceptions as well and thus is intended for use at the top level of elaboration.
-Interrupt and abort exceptions are caught but not logged.
--/
-@[inline] def withLoggingExceptions (x : CommandElabM Unit) : CommandElabCoreM Empty Unit := fun ctx ref =>
-  EIO.catchExceptions (withLogging x ctx ref) (fun _ => pure ())
-
 private def liftAttrM {α} (x : AttrM α) : CommandElabM α := do
   liftCoreM x
 

src/Lean/Elab/Term.lean

@@ -473,6 +473,26 @@ def withoutTacticReuse [Monad m] [MonadWithReaderOf Context m] [MonadOptions m]
       return !cond }
   }) act
 
+@[inherit_doc Core.wrapAsync]
+def wrapAsync (act : Unit → TermElabM α) : TermElabM (EIO Exception α) := do
+  let ctx ← read
+  let st ← get
+  let metaCtx ← readThe Meta.Context
+  let metaSt ← getThe Meta.State
+  Core.wrapAsync fun _ =>
+    act () |>.run ctx |>.run' st |>.run' metaCtx metaSt
+
+@[inherit_doc Core.wrapAsyncAsSnapshot]
+def wrapAsyncAsSnapshot (act : Unit → TermElabM Unit)
+    (desc : String := by exact decl_name%.toString) :
+    TermElabM (BaseIO Language.SnapshotTree) := do
+  let ctx ← read
+  let st ← get
+  let metaCtx ← readThe Meta.Context
+  let metaSt ← getThe Meta.State
+  Core.wrapAsyncAsSnapshot (desc := desc) fun _ =>
+    act () |>.run ctx |>.run' st |>.run' metaCtx metaSt
+
 abbrev TermElabResult (α : Type) := EStateM.Result Exception SavedState α
 
 /--

src/Lean/Language/Basic.lean

@@ -10,9 +10,8 @@ Authors: Sebastian Ullrich
 
 prelude
 import Init.System.Promise
-import Lean.Message
 import Lean.Parser.Types
-import Lean.Elab.InfoTree
+import Lean.Util.Trace
 
 set_option linter.missingDocs true
 
@@ -205,20 +204,6 @@ structure SnapshotTree where
   children : Array (SnapshotTask SnapshotTree)
 deriving Inhabited
 
-/-- Produces trace of given snapshot tree, synchronously waiting on all children. -/
-partial def SnapshotTree.trace (s : SnapshotTree) : CoreM Unit :=
-  go none s
-where go range? s := do
-  let file ← getFileMap
-  let mut desc := f!"{s.element.desc}"
-  if let some range := range? then
-    desc := desc ++ f!"{file.toPosition range.start}-{file.toPosition range.stop} "
-  desc := desc ++ .prefixJoin "\n• " (← s.element.diagnostics.msgLog.toList.mapM (·.toString))
-  if let some t := s.element.infoTree? then
-    trace[Elab.info] (← t.format)
-  withTraceNode `Elab.snapshotTree (fun _ => pure desc) do
-    s.children.toList.forM fun c => go c.range? c.get
-
 /--
   Helper class for projecting a heterogeneous hierarchy of snapshot classes to a homogeneous
   representation. -/
@@ -303,6 +288,14 @@ def SnapshotTree.runAndReport (s : SnapshotTree) (opts : Options) (json := false
 def SnapshotTree.getAll (s : SnapshotTree) : Array Snapshot :=
   s.forM (m := StateM _) (fun s => modify (·.push s)) |>.run #[] |>.2
 
+/-- Returns a task that waits on all snapshots in the tree. -/
+def SnapshotTree.waitAll : SnapshotTree → BaseIO (Task Unit)
+  | mk _ children => go children.toList
+where
+  go : List (SnapshotTask SnapshotTree) → BaseIO (Task Unit)
+    | [] => return .pure ()
+    | t::ts => BaseIO.bindTask t.task fun _ => go ts
+
 /-- Context of an input processing invocation. -/
 structure ProcessingContext extends Parser.InputContext