feat: include async elab tasks in info tree reported to linters and request handlers

Author: Kha

src/Lean/CoreM.lean

@@ -462,11 +462,12 @@ def wrapAsyncAsSnapshot (act : Unit → CoreM Unit) (cancelTk? : Option IO.Cance
   let t ← wrapAsync (cancelTk? := cancelTk?) 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
+      -- reset trace/info state and message log so as not to report them twice
       modify fun st => { st with
         messages := st.messages.markAllReported
         traceState := { tid }
         snapshotTasks := #[]
+        infoState := {}
       }
       try
         withTraceNode `Elab.async (fun _ => return desc) do

src/Lean/Elab/Command.lean

@@ -226,13 +226,16 @@ private def runCore (x : CoreM α) : CommandElabM α := do
   }
   let (ea, coreS) ← liftM x
   modify fun s => { s with
-    env               := coreS.env
-    nextMacroScope    := coreS.nextMacroScope
-    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
+    env                      := coreS.env
+    nextMacroScope           := coreS.nextMacroScope
+    ngen                     := coreS.ngen
+    infoState.trees          := s.infoState.trees.append coreS.infoState.trees
+    -- we assume substitution of `assingment` has already happened, but for lazy assignments we only
+    -- do it at the very end
+    infoState.lazyAssignment := coreS.infoState.lazyAssignment
+    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
 
@@ -299,20 +302,20 @@ Interrupt and abort exceptions are caught but not logged.
   EIO.catchExceptions (withLogging x ctx ref) (fun _ => pure ())
 
 @[inherit_doc Core.wrapAsync]
-def wrapAsync (act : Unit → CommandElabM α) (cancelTk? : Option IO.CancelToken) :
-    CommandElabM (EIO Exception α) := do
+def wrapAsync {α β : Type} (act : α → CommandElabM β) (cancelTk? : Option IO.CancelToken) :
+    CommandElabM (α → EIO Exception β) := do
   let ctx ← read
   let ctx := { ctx with cancelTk? }
   let st ← get
-  return act () |>.run ctx |>.run' st
+  return (act · |>.run ctx |>.run' st)
 
 open Language in
 @[inherit_doc Core.wrapAsyncAsSnapshot]
 -- `CoreM` and `CommandElabM` are too different to meaningfully share this code
-def wrapAsyncAsSnapshot (act : Unit → CommandElabM Unit) (cancelTk? : Option IO.CancelToken)
+def wrapAsyncAsSnapshot {α : Type} (act : α → CommandElabM Unit) (cancelTk? : Option IO.CancelToken)
     (desc : String := by exact decl_name%.toString) :
-    CommandElabM (BaseIO SnapshotTree) := do
-  let t ← wrapAsync (cancelTk? := cancelTk?) fun _ => do
+    CommandElabM (α → BaseIO SnapshotTree) := do
+  let f ← wrapAsync (cancelTk? := cancelTk?) fun a => 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
@@ -323,15 +326,15 @@ def wrapAsyncAsSnapshot (act : Unit → CommandElabM Unit) (cancelTk? : Option I
       }
       try
         withTraceNode `Elab.async (fun _ => return desc) do
-          act ()
+          act a
       catch e =>
         logError e.toMessageData
       finally
         addTraceAsMessages
       get
   let ctx ← read
-  return do
-    match (← t.toBaseIO) with
+  return fun a => do
+    match (← (f a).toBaseIO) with
     | .ok (output, st) =>
       let mut msgs := st.messages
       if !output.isEmpty then
@@ -353,6 +356,7 @@ def wrapAsyncAsSnapshot (act : Unit → CommandElabM Unit) (cancelTk? : Option I
 def logSnapshotTask (task : Language.SnapshotTask Language.SnapshotTree) : CommandElabM Unit :=
   modify fun s => { s with snapshotTasks := s.snapshotTasks.push task }
 
+open Language in
 def runLintersAsync (stx : Syntax) : CommandElabM Unit := do
   if !Elab.async.get (← getOptions) then
     withoutModifyingEnv do
@@ -362,8 +366,13 @@ def runLintersAsync (stx : Syntax) : CommandElabM Unit := do
   -- We only start one task for all linters for now as most linters are fast and we simply want
   -- to unblock elaboration of the next command
   let cancelTk ← IO.CancelToken.new
-  let lintAct ← wrapAsyncAsSnapshot (cancelTk? := cancelTk) fun _ => runLinters stx
-  logSnapshotTask { stx? := none, task := (← BaseIO.asTask lintAct), cancelTk? := cancelTk }
+  let lintAct ← wrapAsyncAsSnapshot (cancelTk? := cancelTk) fun infoSt => do
+    modifyInfoState fun _ => infoSt
+    runLinters stx
+
+  -- linters should have access to the complete info tree
+  let task ← BaseIO.mapTask (t := (← getInfoState).substituteLazy) lintAct
+  logSnapshotTask { stx? := none, task, cancelTk? := cancelTk }
 
 protected def getCurrMacroScope : CommandElabM Nat  := do pure (← read).currMacroScope
 protected def getMainModule     : CommandElabM Name := do pure (← getEnv).mainModule
@@ -579,11 +588,7 @@ def elabCommandTopLevel (stx : Syntax) : CommandElabM Unit := withRef stx do pro
     if let some snap := (← read).snap? then
       snap.new.resolve default
 
-    -- 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
-    for tree in (← getInfoTrees) do
-      trace[Elab.info] (← tree.format)
+    let msgs := (← get).messages
     modify fun st => { st with
       messages := initMsgs ++ msgs
       infoState := { st.infoState with trees := initInfoTrees ++ st.infoState.trees }

src/Lean/Elab/Frontend.lean

@@ -112,9 +112,9 @@ where
         |>.foldl (· ++ ·) {}
       -- In contrast to messages, we should collect info trees only from the top-level command
       -- snapshots as they subsume any info trees reported incrementally by their children.
-      let trees := commands.map (·.finishedSnap.get.infoTree?) |>.filterMap id |>.toPArray'
+      let trees := commands.map (·.infoTreeSnap.get.infoTree?) |>.filterMap id |>.toPArray'
       return {
-        commandState := { snap.finishedSnap.get.cmdState with messages, infoState.trees := trees }
+        commandState := { snap.resultSnap.get.cmdState with messages, infoState.trees := trees }
         parserState := snap.parserState
         cmdPos := snap.parserState.pos
         commands := commands.map (·.stx)

src/Lean/Elab/InfoTree/Main.lean

@@ -5,6 +5,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Wojciech Nawrocki, Leonardo de Moura, Sebastian Ullrich
 -/
 prelude
+import Init.Task
 import Lean.Meta.PPGoal
 import Lean.ReservedNameAction
 
@@ -95,6 +96,13 @@ partial def InfoTree.substitute (tree : InfoTree) (assignment : PersistentHashMa
     | none      => hole id
     | some tree => substitute tree assignment
 
+/-- Applies `s.lazyAssignment` to `s.trees`, asynchronously. -/
+def InfoState.substituteLazy (s : InfoState) : Task InfoState :=
+  Task.mapList (tasks := s.lazyAssignment.toList.map (·.2)) fun _ => { s with
+    trees := s.trees.map (·.substitute <| s.lazyAssignment.map (·.get))
+    lazyAssignment := {}
+  }
+
 /-- Embeds a `CoreM` action in `IO` by supplying the information stored in `info`. -/
 def ContextInfo.runCoreM (info : ContextInfo) (x : CoreM α) : IO α := do
   /-

src/Lean/Elab/InfoTree/Types.lean

@@ -255,6 +255,11 @@ structure InfoState where
   enabled    : Bool := true
   /-- Map from holes in the infotree to child infotrees. -/
   assignment : PersistentHashMap MVarId InfoTree := {}
+  /--
+  Assignments fulfilled by other elaboration tasks. We substitute them only just before reporting
+  the info tree via a snapshot to avoid premature blocking.
+  -/
+  lazyAssignment : PersistentHashMap MVarId (Task InfoTree) := {}
   /-- Pending child trees of a node. -/
   trees      : PersistentArray InfoTree := {}
   deriving Inhabited

src/Lean/Elab/InfoTrees.lean

@@ -17,7 +17,7 @@ def elabInfoTrees : CommandElab
       logError "Info trees are disabled, can not use `#info_trees`."
     else
       elabCommand cmd
-      let infoTrees ← getInfoTrees
+      let infoTrees := (← getInfoState).substituteLazy.get.trees
       for t in infoTrees do
         logInfoAt tk (← t.format)
   | _ => throwUnsupportedSyntax

src/Lean/Language/Lean.lean

@@ -524,12 +524,12 @@ where
       if let some oldNext := old.nextCmdSnap? then do
         let newProm ← IO.Promise.new
         -- can reuse range, syntax unchanged
-        BaseIO.chainTask (sync := true) old.finishedSnap.task fun oldFinished =>
+        BaseIO.chainTask (sync := true) old.resultSnap.task fun oldResult =>
           -- also wait on old command parse snapshot as parsing is cheap and may allow for
           -- elaboration reuse
           BaseIO.chainTask (sync := true) oldNext.task fun oldNext => do
             let cancelTk ← IO.CancelToken.new
-            parseCmd oldNext newParserState oldFinished.cmdState newProm sync cancelTk ctx
+            parseCmd oldNext newParserState oldResult.cmdState newProm sync cancelTk ctx
         prom.resolve <| { old with nextCmdSnap? := some {
           stx? := none
           reportingRange? := some ⟨newParserState.pos, ctx.input.endPos⟩
@@ -582,7 +582,8 @@ where
       prom.resolve <| {
         diagnostics := .empty, stx := .missing, parserState
         elabSnap := default
-        finishedSnap := .finished none { diagnostics := .empty, cmdState }
+        resultSnap := .finished none { diagnostics := .empty, cmdState }
+        infoTreeSnap := .finished none { diagnostics := .empty }
         reportSnap := default
         nextCmdSnap? := none
       }
@@ -592,6 +593,7 @@ where
     let _ ← (if sync then BaseIO.asTask else (.pure <$> ·)) do
       -- definitely resolved in `doElab` task
       let elabPromise ← IO.Promise.new
+      let resultPromise ← IO.Promise.new
       let finishedPromise ← IO.Promise.new
       let reportPromise ← IO.Promise.new
       let minimalSnapshots := internal.cmdlineSnapshots.get cmdState.scopes.head!.opts
@@ -602,11 +604,6 @@ where
       -- report terminal tasks on first line of decl such as not to hide incremental tactics'
       -- progress
       let initRange? := getNiceCommandStartPos? stx |>.map fun pos => ⟨pos, pos⟩
-      let finishedSnap := {
-        stx? := stx'
-        reportingRange? := initRange?
-        task := finishedPromise.result!
-      }
       let next? ← if Parser.isTerminalCommand stx then pure none
         -- for now, wait on "command finished" snapshot before parsing next command
         else some <$> IO.Promise.new
@@ -621,21 +618,55 @@ where
       -- use per-command cancellation token for elaboration so that
       let elabCmdCancelTk ← IO.CancelToken.new
       prom.resolve {
-        diagnostics, finishedSnap, nextCmdSnap?
+        diagnostics, nextCmdSnap?
         stx := stx', parserState := parserState'
         elabSnap := { stx? := stx', task := elabPromise.result!, cancelTk? := some elabCmdCancelTk }
+        resultSnap := { stx? := stx', reportingRange? := initRange?, task := resultPromise.result! }
+        infoTreeSnap := { stx? := stx', reportingRange? := initRange?, task := finishedPromise.result! }
         reportSnap := { stx? := none, reportingRange? := initRange?, task := reportPromise.result! }
       }
       let cmdState ← doElab stx cmdState beginPos
         { old? := old?.map fun old => ⟨old.stx, old.elabSnap⟩, new := elabPromise }
-        finishedPromise elabCmdCancelTk ctx
+        elabCmdCancelTk ctx
+
+      let mut reportedCmdState := cmdState
+      let cmdline := internal.cmdlineSnapshots.get scope.opts && !Parser.isTerminalCommand stx
+      if cmdline then
+        -- discard all metadata apart from the environment; see `internal.cmdlineSnapshots`
+        reportedCmdState := { env := reportedCmdState.env, maxRecDepth := 0 }
+      resultPromise.resolve {
+        diagnostics := (← Snapshot.Diagnostics.ofMessageLog cmdState.messages)
+        traces := cmdState.traceState
+        cmdState := reportedCmdState
+      }
+
+      -- report info tree when relevant tasks are finished
+      BaseIO.chainTask (sync := true) (t := cmdState.infoState.substituteLazy) fun infoSt => do
+        let infoTree := infoSt.trees[0]!
+        let opts := cmdState.scopes.head!.opts
+        let mut msgLog := MessageLog.empty
+        if (← isTracingEnabledForCore `Elab.info opts) then
+          if let .ok msg ← infoTree.format.toBaseIO then
+            let data := .tagged `trace <| .trace { cls := `Elab.info } .nil #[msg]
+            msgLog := msgLog.add {
+              fileName := ctx.fileName
+              severity := MessageSeverity.information
+              pos      := ctx.fileMap.toPosition beginPos
+              data     := data
+            }
+        finishedPromise.resolve {
+          diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+          infoTree? := infoTree
+        }
+
+      -- report traces when *all* tasks are finished
       let traceTask ←
         if (← isTracingEnabledForCore `Elab.snapshotTree cmdState.scopes.head!.opts) then
           -- We want to trace all of `CommandParsedSnapshot` but `traceTask` is part of it, so let's
           -- create a temporary snapshot tree containing all tasks but it
           let snaps := #[
             { stx? := stx', task := elabPromise.result!.map (sync := true) toSnapshotTree },
-            { stx? := stx', task := finishedPromise.result!.map (sync := true) toSnapshotTree }] ++
+            { stx? := stx', task := resultPromise.result!.map (sync := true) toSnapshotTree }] ++
             cmdState.snapshotTasks
           let tree := SnapshotTree.mk { diagnostics := .empty } snaps
           BaseIO.bindTask (← tree.waitAll) fun _ => do
@@ -665,8 +696,8 @@ where
         parseCmd none parserState cmdState next (sync := false) elabCmdCancelTk ctx
 
   doElab (stx : Syntax) (cmdState : Command.State) (beginPos : String.Pos)
-      (snap : SnapshotBundle DynamicSnapshot) (finishedPromise : IO.Promise CommandFinishedSnapshot)
-      (cancelTk : IO.CancelToken) : LeanProcessingM Command.State := do
+      (snap : SnapshotBundle DynamicSnapshot) (cancelTk : IO.CancelToken) :
+      LeanProcessingM Command.State := do
     let ctx ← read
     let scope := cmdState.scopes.head!
     -- reset per-command state
@@ -693,21 +724,9 @@ where
         data     := output
       }
     let cmdState : Command.State := { cmdState with messages }
-    let mut reportedCmdState := cmdState
     -- definitely resolve eventually
     snap.new.resolve <| .ofTyped { diagnostics := .empty : SnapshotLeaf }
 
-    let infoTree : InfoTree := cmdState.infoState.trees[0]!
-    let cmdline := internal.cmdlineSnapshots.get scope.opts && !Parser.isTerminalCommand stx
-    if cmdline then
-      -- discard all metadata apart from the environment; see `internal.cmdlineSnapshots`
-      reportedCmdState := { env := reportedCmdState.env, maxRecDepth := 0 }
-    finishedPromise.resolve {
-      diagnostics := (← Snapshot.Diagnostics.ofMessageLog cmdState.messages)
-      infoTree? := infoTree
-      traces := cmdState.traceState
-      cmdState := reportedCmdState
-    }
     -- The reported `cmdState` in the snapshot may be minimized as seen above, so we return the full
     -- state here for further processing on the same thread
     return cmdState
@@ -735,6 +754,6 @@ where goCmd snap :=
   if let some next := snap.nextCmdSnap? then
     goCmd next.get
   else
-    snap.finishedSnap.get.cmdState
+    snap.resultSnap.get.cmdState
 
 end Lean

src/Lean/Language/Lean/Types.lean

@@ -25,12 +25,15 @@ private def pushOpt (a? : Option α) (as : Array α) : Array α :=
 
 /-! The hierarchy of Lean snapshot types -/
 
-/-- Snapshot after elaboration of the entire command. -/
-structure CommandFinishedSnapshot extends Language.Snapshot where
+/--
+Snapshot after command elaborator has returned. Also contains diagnostics from the elaborator's main
+task. Asynchronous elaboration tasks may not yet be finished.
+-/
+structure CommandResultSnapshot extends Language.Snapshot where
   /-- Resulting elaboration state. -/
   cmdState : Command.State
 deriving Nonempty
-instance : ToSnapshotTree CommandFinishedSnapshot where
+instance : ToSnapshotTree CommandResultSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[]⟩
 
 /-- State after a command has been parsed. -/
@@ -44,8 +47,12 @@ structure CommandParsedSnapshot extends Snapshot where
   elaborator.
    -/
   elabSnap : SnapshotTask DynamicSnapshot
-  /-- State after processing is finished. -/
-  finishedSnap : SnapshotTask CommandFinishedSnapshot
+  /-- State after command elaborator has returned. -/
+  resultSnap : SnapshotTask CommandResultSnapshot
+  /--
+  State after all elaborator tasks are finished. In particular, contains the complete info tree.
+  -/
+  infoTreeSnap : SnapshotTask SnapshotLeaf
   /-- Additional, untyped snapshots used for reporting, not reuse. -/
   reportSnap : SnapshotTask SnapshotTree
   /-- Next command, unless this is a terminal command. -/
@@ -55,7 +62,8 @@ partial instance : ToSnapshotTree CommandParsedSnapshot where
   toSnapshotTree := go where
     go s := ⟨s.toSnapshot,
       #[s.elabSnap.map (sync := true) toSnapshotTree,
-        s.finishedSnap.map (sync := true) toSnapshotTree,
+        s.resultSnap.map (sync := true) toSnapshotTree,
+        s.infoTreeSnap.map (sync := true) toSnapshotTree,
         s.reportSnap] |>
         pushOpt (s.nextCmdSnap?.map (·.map (sync := true) go))⟩
 

src/Lean/Server/FileWorker/Utils.lean

@@ -28,11 +28,11 @@ private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     } <| .delayed <| headerSuccess.firstCmdSnap.task.asServerTask.bindCheap go
 where
   go cmdParsed :=
-    cmdParsed.finishedSnap.task.asServerTask.mapCheap fun finished =>
+    cmdParsed.resultSnap.task.asServerTask.mapCheap fun result =>
       .ok <| .cons {
         stx := cmdParsed.stx
         mpState := cmdParsed.parserState
-        cmdState := finished.cmdState
+        cmdState := result.cmdState
       } (match cmdParsed.nextCmdSnap? with
         | some next => .delayed <| next.task.asServerTask.bindCheap go
         | none => .nil)

src/Lean/Server/Requests.lean

@@ -304,10 +304,9 @@ def findCmdDataAtPos
   findCmdParsedSnap doc hoverPos |>.bindCheap fun
     | some cmdParsed => toSnapshotTree cmdParsed |>.findInfoTreeAtPos doc.meta.text hoverPos includeStop |>.bindCheap fun
       | some infoTree => .pure <| some (cmdParsed.stx, infoTree)
-      | none          => cmdParsed.finishedSnap.task.asServerTask.mapCheap fun s =>
-        -- the parser returns exactly one command per snapshot, and the elaborator creates exactly one node per command
-        assert! s.cmdState.infoState.trees.size == 1
-        some (cmdParsed.stx, s.cmdState.infoState.trees[0]!)
+      | none          => cmdParsed.infoTreeSnap.task.asServerTask.mapCheap fun s =>
+        assert! s.infoTree?.isSome
+        some (cmdParsed.stx, s.infoTree?.get!)
     | none => .pure none
 
 open Language in