feat: parallel progress notifications

src/Lean/Data/Lsp/Basic.lean

@@ -32,7 +32,7 @@ offsets. For diagnostics, one-based `Lean.Position`s are used internally.
 structure Position where
   line : Nat
   character : Nat
-  deriving Inhabited, BEq, Ord, Hashable, ToJson, FromJson
+  deriving Inhabited, BEq, Ord, Hashable, ToJson, FromJson, Repr
 
 instance : ToString Position := ⟨fun p =>
   "(" ++ toString p.line ++ ", " ++ toString p.character ++ ")"⟩
@@ -43,7 +43,7 @@ instance : LE Position := leOfOrd
 structure Range where
   start : Position
   «end» : Position
-  deriving Inhabited, BEq, Hashable, ToJson, FromJson, Ord
+  deriving Inhabited, BEq, Hashable, ToJson, FromJson, Ord, Repr
 
 instance : LT Range := ltOfOrd
 instance : LE Range := leOfOrd

src/Lean/Elab/MutualDef.lean

@@ -215,6 +215,14 @@ private def elabHeaders (views : Array DefView)
             return newHeader
       if let some snap := view.headerSnap? then
         let (tacStx?, newTacTask?) ← mkTacTask view.value tacPromise
+        let bodySnap :=
+          -- Only use first line of body as range when we have incremental tactics as otherwise we
+          -- would cover their progress
+          { range? := if newTacTask?.isSome then
+              view.ref.getPos?.map fun pos => ⟨pos, pos⟩
+            else
+              getBodyTerm? view.value |>.getD view.value |>.getRange?
+            task := bodyPromise.result }
         snap.new.resolve <| some {
           diagnostics :=
             (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getAndEmptyMessageLog))
@@ -223,7 +231,7 @@ private def elabHeaders (views : Array DefView)
           tacStx?
           tacSnap? := newTacTask?
           bodyStx := view.value
-          bodySnap := mkBodyTask view.value bodyPromise
+          bodySnap
         }
         newHeader := { newHeader with
           -- We should only forward the promise if we are actually waiting on the
@@ -245,12 +253,6 @@ where
     guard whereDeclsOpt.isNone
     return body
 
-  /-- Creates snapshot task with appropriate range from body syntax and promise. -/
-  mkBodyTask (body : Syntax) (new : IO.Promise (Option BodyProcessedSnapshot)) :
-      Language.SnapshotTask (Option BodyProcessedSnapshot) :=
-    let rangeStx := getBodyTerm? body |>.getD body
-    { range? := rangeStx.getRange?, task := new.result }
-
   /--
   If `body` allows for incremental tactic reporting and reuse, creates a snapshot task out of the
   passed promise with appropriate range, otherwise immediately resolves the promise to a dummy

src/Lean/Language/Lean.lean

@@ -292,6 +292,13 @@ If the option is defined in this library, use '-D{`weak ++ name}' to set it cond
 
   return opts
 
+private def getNiceCommandStartPos? (stx : Syntax) : Option String.Pos := do
+  let mut stx := stx
+  if stx[0].isOfKind ``Command.declModifiers then
+    -- modifiers are morally before the actual declaration
+    stx := stx[1]
+  stx.getPos?
+
 /--
 Entry point of the Lean language processor.
 
@@ -540,15 +547,10 @@ where
       let elabPromise ← IO.Promise.new
       let finishedPromise ← IO.Promise.new
       let reportPromise ← IO.Promise.new
-      -- (Try to) use last line of command as range for final snapshot task. This ensures we do not
-      -- retract the progress bar to a previous position in case the command support incremental
-      -- reporting but has significant work after resolving its last incremental promise, such as
-      -- final type checking; if it does not support incrementality, `elabSnap` constructed in
-      -- `parseCmd` and containing the entire range of the command will determine the reported
-      -- progress and be resolved effectively at the same time as this snapshot task, so `tailPos` is
-      -- irrelevant in this case.
-      let endRange? := stx.getTailPos?.map fun pos => ⟨pos, pos⟩
-      let finishedSnap := { range? := endRange?, task := finishedPromise.result }
+      -- 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 := { range? := initRange?, task := finishedPromise.result }
       let tacticCache ← old?.map (·.tacticCache) |>.getDM (IO.mkRef {})
 
       let minimalSnapshots := internal.cmdlineSnapshots.get cmdState.scopes.head!.opts
@@ -566,7 +568,7 @@ where
         diagnostics, finishedSnap, tacticCache, nextCmdSnap?
         stx := stx', parserState := parserState'
         elabSnap := { range? := stx.getRange?, task := elabPromise.result }
-        reportSnap := { range? := endRange?, task := reportPromise.result }
+        reportSnap := { range? := initRange?, task := reportPromise.result }
       }
       let cmdState ← doElab stx cmdState beginPos
         { old? := old?.map fun old => ⟨old.stx, old.elabSnap⟩, new := elabPromise }
@@ -597,7 +599,7 @@ where
           pure <| .pure <| .mk { diagnostics := .empty } #[]
       reportPromise.resolve <|
         .mk { diagnostics := .empty } <|
-          cmdState.snapshotTasks.push { range? := endRange?, task := traceTask }
+          cmdState.snapshotTasks.push { range? := initRange?, task := traceTask }
       if let some next := next? then
         -- We're definitely off the fast-forwarding path now
         parseCmd none parserState cmdState next (sync := false) ctx

src/Lean/Server/FileWorker.lean

@@ -125,11 +125,6 @@ section Elab
     newInfoTrees : Array Elab.InfoTree := #[]
     /-- Whether we encountered any snapshot with `Snapshot.isFatal`. -/
     hasFatal := false
-    /--
-    Last `Snapshot.range?` encountered that was not `none`, if any. We use this as a fallback when
-    reporting progress as we should always report *some* range when waiting on a task.
-    -/
-    lastRange? : Option String.Range := none
   deriving Inhabited
 
   register_builtin_option server.reportDelayMs : Nat := {
@@ -172,34 +167,73 @@ This option can only be set on the command line, not in the lakefile or via `set
     See also section "Communication" in Lean/Server/README.md.
 
     Debouncing: we only report information
-    * after first waiting for `reportDelayMs`, to give trivial tasks a chance to finish
-    * when first blocking, i.e. not before skipping over any unchanged snapshots and such trivial
-      tasks
-    * afterwards, each time new information is found in a snapshot
-    * at the very end, if we never blocked (e.g. emptying a file should make
-      sure to empty diagnostics as well eventually) -/
+    1. after first waiting for `reportDelayMs`, to give trivial tasks a chance to finish
+    2. when first blocking, i.e. not before skipping over any unchanged snapshots and such trivial
+       tasks
+    3. afterwards, each time new information is found in a snapshot
+    4. at the very end, if we never blocked (e.g. emptying a file should make
+       sure to empty diagnostics as well eventually) -/
   private partial def reportSnapshots (ctx : WorkerContext) (doc : EditableDocumentCore)
       (cancelTk : CancelToken) : BaseIO (Task Unit) := do
     let t ← BaseIO.asTask do
-      IO.sleep (server.reportDelayMs.get ctx.cmdlineOpts).toUInt32
+      IO.sleep (server.reportDelayMs.get ctx.cmdlineOpts).toUInt32  -- "Debouncing 1."
     BaseIO.bindTask t fun _ => do
-      BaseIO.bindTask (← go (toSnapshotTree doc.initSnap) {}) fun st => do
-        if (← cancelTk.isSet) then
-          return .pure ()
-
-        -- callback at the end of reporting
-        if st.hasFatal then
-          ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta 0 .fatalError
-        else
-          ctx.chanOut.send <| mkFileProgressDoneNotification doc.meta
-        unless st.hasBlocked do
-          publishDiagnostics ctx doc
-        -- This will overwrite existing ilean info for the file, in case something
-        -- went wrong during the incremental updates.
-        ctx.chanOut.send (← mkIleanInfoFinalNotification doc.meta st.allInfoTrees)
+      let (_, st) ← handleTasks #[.pure <| toSnapshotTree doc.initSnap] |>.run {}
+      if (← cancelTk.isSet) then
         return .pure ()
+
+      -- callback at the end of reporting
+      if st.hasFatal then
+        ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta 0 .fatalError
+      else
+        ctx.chanOut.send <| mkFileProgressDoneNotification doc.meta
+      unless st.hasBlocked do  -- "Debouncing 4."
+        publishDiagnostics ctx doc
+      -- This will overwrite existing ilean info for the file, in case something
+      -- went wrong during the incremental updates.
+      ctx.chanOut.send (← mkIleanInfoFinalNotification doc.meta st.allInfoTrees)
+      return .pure ()
   where
-    go (node : SnapshotTree) (st : ReportSnapshotsState) : BaseIO (Task ReportSnapshotsState) := do
+    /--
+    Given an array of possibly-unfinished tasks, handles them, possibly after waiting for one of
+    them to finish.
+    -/
+    handleTasks (ts : Array (SnapshotTask SnapshotTree)) : StateT ReportSnapshotsState BaseIO Unit := do
+      let ts ← ts.flatMapM handleFinished
+      -- all `ts` are now (likely) in-progress, report them
+      sendFileProgress ts
+      -- check again whether this has changed before commiting to waiting
+      if (← ts.anyM (IO.hasFinished ·.task)) then
+        handleTasks ts
+      else if h : ts.size > 0 then
+        if !(← get).hasBlocked then  -- "Debouncing 2."
+          publishDiagnostics ctx doc
+          modify fun st => { st with hasBlocked := true }
+        -- wait for at least one task to finish; there is a race condition here where a task may
+        -- have finished between the previous check and this line but we accept the progress
+        -- notifications being temporarily out of date in this case
+        let _ ← IO.waitAny (ts.map (·.task) |>.toList)
+          (by simp only [Array.toList_map, List.length_map, Array.length_toList, gt_iff_lt, h])
+        handleTasks ts
+
+    /-- Recursively handles finished tasks and replaces them with their unfinished children. -/
+    handleFinished (t : SnapshotTask SnapshotTree) :
+        StateT ReportSnapshotsState BaseIO (Array (SnapshotTask SnapshotTree)) := do
+      if (← IO.hasFinished t.task) then
+        handleNode t.task.get
+        let ts ← t.task.get.children.flatMapM handleFinished
+        -- limit children's reported range to that of the parent, if any, to avoid strange
+        -- non-monotonic progress updates; replace missing children's ranges with parent's
+        return ts.map (fun t' => { t' with range? := match t.range?, t'.range? with
+          | some r, some r' =>
+            let (_, _) := (minOfLe (α := String.Pos), maxOfLe (α := String.Pos))
+            some { start := max r.start r'.start, stop := min r.stop r'.stop }
+          | r?, r?' => r?' <|> r? })
+      else
+        return #[t]
+
+    /-- Handles information of a single now-finished snapshot. -/
+    handleNode (node : SnapshotTree) : StateT ReportSnapshotsState BaseIO Unit := do
       if node.element.diagnostics.msgLog.hasUnreported then
         let diags ←
           if let some memorized ← node.element.diagnostics.interactiveDiagsRef?.bindM fun ref => do
@@ -212,35 +246,29 @@ This option can only be set on the command line, not in the lakefile or via `set
               cacheRef.set <| some <| .mk { diags : MemorizedInteractiveDiagnostics }
             pure diags
         doc.diagnosticsRef.modify (· ++ diags)
-        if st.hasBlocked then
+        if (← get).hasBlocked then
           publishDiagnostics ctx doc
 
-      let mut st := { st with hasFatal := st.hasFatal || node.element.isFatal }
+      modify fun st => { st with hasFatal := st.hasFatal || node.element.isFatal }
 
       if let some itree := node.element.infoTree? then
-        let mut newInfoTrees := st.newInfoTrees.push itree
-        if st.hasBlocked then
+        let mut newInfoTrees := (← get).newInfoTrees.push itree
+        if (← get).hasBlocked then
           ctx.chanOut.send (← mkIleanInfoUpdateNotification doc.meta newInfoTrees)
           newInfoTrees := #[]
-        st := { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
-
-      goSeq st node.children.toList
-
-    goSeq (st : ReportSnapshotsState) :
-        List (SnapshotTask SnapshotTree) → BaseIO (Task ReportSnapshotsState)
-      | [] => return .pure st
-      | t::ts => do
-        let mut st := st
-        st := { st with lastRange? := t.range? <|> st.lastRange? }
-        unless (← IO.hasFinished t.task) do
-          -- report *some* recent range even if `t.range?` is `none`; see also `State.lastRange?`
-          if let some range := st.lastRange? then
-            ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta range.start
-          if !st.hasBlocked then
-            publishDiagnostics ctx doc
-            st := { st with hasBlocked := true }
-        BaseIO.bindTask t.task fun t => do
-          BaseIO.bindTask (← go t st) (goSeq · ts)
+        modify fun st => { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
+
+    /-- Reports given tasks' ranges, merging overlapping ones. -/
+    sendFileProgress (tasks : Array (SnapshotTask SnapshotTree)) : StateT ReportSnapshotsState BaseIO Unit := do
+      let ranges := tasks.filterMap (·.range?)
+      let ranges := ranges.qsort (·.start < ·.start)
+      let ranges := ranges.foldl (init := #[]) fun rs r => match rs[rs.size - 1]? with
+        | some last => if last.stop < r.start then rs.push r else rs.pop.push { last with stop := r.stop }
+        | none => rs.push r
+      let ranges := ranges.map (·.toLspRange doc.meta.text)
+      let notifs := ranges.map ({ range := ·, kind := .processing })
+      ctx.chanOut.send <| mkFileProgressNotification doc.meta notifs
+
 end Elab
 
 -- Pending requests are tracked so they can be canceled