feat: inlay hint refinements (#6959)

This PR implements a number of refinements for the auto-implicit inlay
hints implemented in #6768.
Specifically:
- In #6768, there was a bug where the inlay hint edit delay could
accumulate on successive edits, which meant that it could sometimes take
much longer for inlay hints to show up. This PR implements the basic
infrastructure for request cancellation and implements request
cancellation for semantic tokens and inlay hints to resolve the issue.
With this edit delay bug fixed, it made more sense to increase the edit
delay slightly from 2000ms to 3000ms.
- In #6768, we applied the edit delay to every single inlay hint request
in order to reduce the amount of inlay hint flickering. This meant that
the edit delay also had a significant effect on how far inlay hints
would lag behind the file progress bar. This PR adjusts the edit delay
logic so that it only affects requests sent directly after a
corresponding `didChange` notification. Once the edit delay is used up,
all further semantic token requests are responded to without delay, so
that the only latency that affects how far the inlay hints lag behind
the progress bar is how often we emit refresh requests and how long VS
Code takes to respond to them.
- For inlay hints, refresh requests are now emitted 500ms after a
response to an inlay hint request, not 2000ms, which means that after
the edit delay, inlay hints should only lag behind the progress bar by
about up to 500ms. This is justifiable for inlay hints because the
response should be much smaller than e.g. is the case for semantic
tokens.
- In #6768, 'Restart File' did not prompt a refresh, but it does now.
- VS Code does not immediately remove old inlay hints from the document
when they are applied. In #6768, this meant that inlay hints would
linger around for a bit once applied. To mitigate this issue, this PR
adjusts the inlay hint edit delay logic to identify edits sent from the
client as being inlay hint applications, and sets the edit delay to 0ms
for the inlay hint requests following it. This means that inlay hints
are now applied immediately.
- In #6768, hovering over single-letter auto-implicit inlay hints was a
bit finicky because VS Code uses the regular cursor icon on inlay hints,
not the thin text cursor icon, which means that it is easy to put the
cursor in the wrong spot. We now add the separation character (` ` or
`{`) preceding an auto-implicit to the hover range as well, which makes
hovering over inlay hints much smoother.

src/Lean/Elab/InfoTree/InlayHints.lean

@@ -30,6 +30,7 @@ inductive InlayHintKind where
 structure InlayHintTextEdit where
   range   : String.Range
   newText : String
+  deriving BEq
 
 structure InlayHintInfo where
   position     : String.Pos

src/Lean/Elab/Term.lean

@@ -2005,14 +2005,24 @@ def addAutoBoundImplicitsInlayHint (autos : Array Expr) (inlayHintPos : String.P
   let deferredResolution ih := do
     let .parts ps := ih.label
       | return ih
-    let ps ← ps.mapIdxM fun i p => do
+    let mut ps' := #[]
+    for h : i in [:ps.size] do
+      let p := ps[i]
       let some (part, some auto) := labelParts[i]?
-        | return p
+        | ps' := ps'.push p
+          continue
       let type := toString <| ← Meta.ppExpr <| ← instantiateMVars (← inferType auto)
-      return { p with
-        tooltip? := s!"{part.value} : {type}"
-      }
-    return { ih with label := .parts ps }
+      let tooltip := s!"{part.value} : {type}"
+      ps' := ps'.push { p with tooltip? := tooltip }
+      let some separatorPart := ps'[ps'.size - 2]?
+        | continue
+      -- We assign the leading `{` and the separation spaces the same tooltip as the auto-implicit
+      -- following it. The reason for this is that VS Code does not display a text cursor
+      -- on auto-implicits, but a regular cursor, and hitting single character auto-implicits
+      -- with that cursor can be a bit tricky. Adding the leading space or the opening `{` to the
+      -- tooltip area makes this much easier.
+      ps' := ps'.set! (ps'.size - 2) { separatorPart with tooltip? := tooltip }
+    return { ih with label := .parts ps' }
   pushInfoLeaf <| .ofCustomInfo {
       position := inlayHintPos
       label := .parts <| labelParts.map (·.1)

src/Lean/Server/AsyncList.lean

@@ -6,6 +6,7 @@ Authors: Wojciech Nawrocki
 -/
 prelude
 import Init.System.IO
+import Init.System.Promise
 
 namespace IO
 
@@ -114,34 +115,58 @@ partial def getFinishedPrefix : AsyncList ε α → BaseIO (List α × Option ε
       | Except.error e => pure ⟨[], some e, true⟩
     else pure ⟨[], none, false⟩
 
-partial def getFinishedPrefixWithTimeout (xs : AsyncList ε α) (timeoutMs : UInt32) : BaseIO (List α × Option ε × Bool) := do
-  let timeoutTask : Task (Unit ⊕ Except ε (AsyncList ε α)) ← BaseIO.asTask (prio := .dedicated) do
-    IO.sleep timeoutMs
-    return .inl ()
+partial def getFinishedPrefixWithTimeout (xs : AsyncList ε α) (timeoutMs : UInt32)
+    (cancelTk? : Option (Task Unit) := none) : BaseIO (List α × Option ε × Bool) := do
+  let timeoutTask : Task (Unit ⊕ Except ε (AsyncList ε α)) ←
+    if timeoutMs == 0 then
+      pure <| Task.pure (Sum.inl ())
+    else
+      BaseIO.asTask (prio := .dedicated) do
+        IO.sleep timeoutMs
+        return .inl ()
   go timeoutTask xs
 where
-  go (timeoutTask : Task (Unit ⊕ Except ε (AsyncList ε α))) (xs : AsyncList ε α) : BaseIO (List α × Option ε × Bool) := do
+  go (timeoutTask : Task (Unit ⊕ Except ε (AsyncList ε α)))
+      (xs : AsyncList ε α) : BaseIO (List α × Option ε × Bool) := do
     match xs with
     | cons hd tl =>
       let ⟨tl, e?, isComplete⟩ ← go timeoutTask tl
       return ⟨hd :: tl, e?, isComplete⟩
     | nil => return ⟨[], none, true⟩
     | delayed tl =>
-      let r ← IO.waitAny [
-        timeoutTask,
-        tl.map (sync := true) .inr
-      ]
+      let tl := tl.map (sync := true) .inr
+      let cancelTk? := do return (← cancelTk?).map (sync := true) .inl
+      let tasks : { t : List _ // t.length > 0 } :=
+        match cancelTk? with
+        | none => ⟨[tl, timeoutTask], by exact Nat.zero_lt_succ _⟩
+        | some cancelTk => ⟨[cancelTk, tl, timeoutTask], by exact Nat.zero_lt_succ _⟩
+      let r ← IO.waitAny tasks.val (h := tasks.property)
       match r with
-      | .inl _ => return ⟨[], none, false⟩ -- Timeout - stop waiting
+      | .inl _ => return ⟨[], none, false⟩ -- Timeout or cancellation - stop waiting
       | .inr (.ok tl) => go timeoutTask tl
       | .inr (.error e) => return ⟨[], some e, true⟩
 
-partial def getFinishedPrefixWithConsistentLatency (xs : AsyncList ε α) (latencyMs : UInt32) : BaseIO (List α × Option ε × Bool) := do
+partial def getFinishedPrefixWithConsistentLatency (xs : AsyncList ε α) (latencyMs : UInt32)
+    (cancelTk? : Option (Task Unit) := none) : BaseIO (List α × Option ε × Bool) := do
   let timestamp ← IO.monoMsNow
-  let r ← xs.getFinishedPrefixWithTimeout latencyMs
+  let r ← xs.getFinishedPrefixWithTimeout latencyMs cancelTk?
   let passedTimeMs := (← IO.monoMsNow) - timestamp
-  IO.sleep <| (latencyMs.toNat - passedTimeMs).toUInt32
+  let remainingLatencyMs := (latencyMs.toNat - passedTimeMs).toUInt32
+  sleepWithCancellation remainingLatencyMs
   return r
+where
+  sleepWithCancellation (sleepDurationMs : UInt32) : BaseIO Unit := do
+    if sleepDurationMs == 0 then
+      return
+    let some cancelTk := cancelTk?
+      | IO.sleep sleepDurationMs
+        return
+    if ← IO.hasFinished cancelTk then
+      return
+    let sleepTask ← BaseIO.asTask (prio := .dedicated) do
+      IO.sleep sleepDurationMs
+    IO.waitAny [sleepTask, cancelTk]
+
 
 def waitHead? (as : AsyncList ε α) : Task (Except ε (Option α)) :=
   as.waitFind? fun _ => true

src/Lean/Server/FileWorker.lean

@@ -266,8 +266,12 @@ This option can only be set on the command line, not in the lakefile or via `set
           BaseIO.bindTask (← go t st) (goSeq · ts)
 end Elab
 
+structure PendingRequest where
+  requestTask : Task (Except IO.Error Unit)
+  cancelTk    : RequestCancellationToken
+
 -- Pending requests are tracked so they can be canceled
-abbrev PendingRequestMap := RBMap RequestID (Task (Except IO.Error Unit)) compare
+abbrev PendingRequestMap := RBMap RequestID PendingRequest compare
 
 structure AvailableImportsCache where
   availableImports       : ImportCompletion.AvailableImports
@@ -367,12 +371,14 @@ section Initialization
     let stickyDiagnosticsRef ← IO.mkRef ∅
     let chanOut ← mkLspOutputChannel maxDocVersionRef
     let srcSearchPathPromise ← IO.Promise.new
+    let timestamp ← IO.monoMsNow
     let partialHandlersRef ← IO.mkRef <| RBMap.fromArray (cmp := compare) <|
-      (← partialLspRequestHandlerMethods).map fun (method, refreshMethod) =>
+      (← partialLspRequestHandlerMethods).map fun (method, refreshMethod, _) =>
         (method, {
           refreshMethod
           requestsInFlight := 0
-          pendingRefreshInfo? := none
+          -- Emit a refresh request after a file worker restart.
+          pendingRefreshInfo? := some { lastRefreshTimestamp := timestamp, successiveRefreshAttempts := 0 }
         })
     let processor := Language.Lean.process (setupImports meta opts chanOut srcSearchPathPromise)
     let processor ← Language.mkIncrementalProcessor processor
@@ -485,21 +491,30 @@ section NotificationHandling
     let ctx ← read
     let st ← get
     let oldDoc := (←get).doc
+    let cancelTk ← RequestCancellationToken.new
     let newVersion := docId.version?.getD 0
     let _ ← IO.mapTask (t := st.srcSearchPathTask) fun srcSearchPath =>
       let rc : RequestContext :=
         { rpcSessions := st.rpcSessions
           srcSearchPath
           doc := oldDoc
+          cancelTk
           hLog := ctx.hLog
           initParams := ctx.initParams }
       RequestM.runInIO (handleOnDidChange p) rc
     if ¬ changes.isEmpty then
       let newDocText := foldDocumentChanges changes oldDoc.meta.text
       updateDocument ⟨docId.uri, newVersion, newDocText, oldDoc.meta.dependencyBuildMode⟩
+      for (_, r) in st.pendingRequests do
+        r.cancelTk.cancel .edit
+
 
   def handleCancelRequest (p : CancelParams) : WorkerM Unit := do
-    updatePendingRequests (fun pendingRequests => pendingRequests.erase p.id)
+    let st ← get
+    let some r := st.pendingRequests.find? p.id
+      | return
+    r.cancelTk.cancel .cancelRequest
+    set <| { st with pendingRequests := st.pendingRequests.erase p.id }
 
   /--
   Received from the watchdog when a dependency of this file is detected as being stale.
@@ -572,9 +587,9 @@ section MessageHandling
     | "$/lean/rpc/keepAlive"   => handle RpcKeepAliveParams handleRpcKeepAlive
     | _                        => throwServerError s!"Got unsupported notification method: {method}"
 
-  def queueRequest (id : RequestID) (requestTask : Task (Except IO.Error Unit))
+  def queueRequest (id : RequestID) (r : PendingRequest)
       : WorkerM Unit := do
-    updatePendingRequests (fun pendingRequests => pendingRequests.insert id requestTask)
+    updatePendingRequests (·.insert id r)
 
   open Widget RequestM Language in
   def handleGetInteractiveDiagnosticsRequest (params : GetInteractiveDiagnosticsParams) :
@@ -671,35 +686,42 @@ section MessageHandling
       ctx.chanOut.send <| .responseError id .internalError (toString e) none
       return
 
+    let cancelTk ← RequestCancellationToken.new
     -- we assume that any other request requires at least the search path
     -- TODO: move into language-specific request handling
-    let t ← IO.bindTask st.srcSearchPathTask fun srcSearchPath => do
+    let requestTask ← IO.bindTask st.srcSearchPathTask fun srcSearchPath => do
       let rc : RequestContext :=
         { rpcSessions := st.rpcSessions
           srcSearchPath
           doc := st.doc
+          cancelTk
           hLog := ctx.hLog
           initParams := ctx.initParams }
       let t? ← EIO.toIO' <| handleLspRequest method params rc
       let t₁ ← match t? with
         | Except.error e =>
-          IO.asTask do
-            ctx.chanOut.send <| e.toLspResponseError id
+            emitResponse ctx (isComplete := false) <| e.toLspResponseError id
+            pure <| Task.pure <| .ok ()
         | Except.ok t => (IO.mapTask · t) fun
           | Except.ok r => do
-            ctx.chanOut.send <| .response id (toJson r.response)
-            let timestamp ← IO.monoMsNow
-            ctx.modifyPartialHandler method fun h => { h with
-              requestsInFlight := h.requestsInFlight - 1
-              pendingRefreshInfo? :=
-                if r.isComplete then
-                  none
-                else
-                  some { lastRefreshTimestamp := timestamp, successiveRefreshAttempts := 0 }
-            }
+            emitResponse ctx (isComplete := r.isComplete) <| .response id (toJson r.response)
           | Except.error e =>
-            ctx.chanOut.send <| e.toLspResponseError id
-    queueRequest id t
+            emitResponse ctx (isComplete := false) <| e.toLspResponseError id
+    queueRequest id { cancelTk, requestTask }
+
+  where
+
+    emitResponse (ctx : WorkerContext) (m : JsonRpc.Message) (isComplete : Bool) : IO Unit := do
+      ctx.chanOut.send m
+      let timestamp ← IO.monoMsNow
+      ctx.modifyPartialHandler method fun h => { h with
+        requestsInFlight := h.requestsInFlight - 1
+        pendingRefreshInfo? :=
+          if isComplete then
+            none
+          else
+            some { lastRefreshTimestamp := timestamp, successiveRefreshAttempts := 0 }
+      }
 
   def handleResponse (_ : RequestID) (_ : Json) : WorkerM Unit :=
     return -- The only response that we currently expect here is always empty
@@ -720,7 +742,7 @@ section MainLoop
           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
+    let pendingRequests ← st.pendingRequests.foldM (fun acc id r => filterFinishedTasks acc id r.requestTask) st.pendingRequests
     st := { st with pendingRequests }
 
     -- Opportunistically (i.e. when we wake up on messages) check if any RPC session has expired.
@@ -750,15 +772,14 @@ section MainLoop
 end MainLoop
 
 def runRefreshTasks : WorkerM (Array (Task Unit)) := do
-  let timeUntilRefreshMs := 2000
   -- We limit the amount of successive refresh attempts in case the user has switched files,
   -- in which case VS Code won't respond to any refresh request for the given file.
   -- Since we don't want to spam the client with refresh requests for every single file that they
   -- switched away from, we limit the amount of attempts.
   let maxSuccessiveRefreshAttempts := 10
   let ctx ← read
   let mut tasks := #[]
-  for (method, refreshMethod) in ← partialLspRequestHandlerMethods do
+  for (method, refreshMethod, refreshIntervalMs) in ← partialLspRequestHandlerMethods do
     tasks := tasks.push <| ← BaseIO.asTask (prio := .dedicated) do
       while true do
         let lastRefreshTimestamp? ← ctx.modifyGetPartialHandler method fun h => Id.run do
@@ -768,14 +789,14 @@ def runRefreshTasks : WorkerM (Array (Task Unit)) := do
             return (none, { h with pendingRefreshInfo? := none })
           return (some info.lastRefreshTimestamp, h)
         let some lastRefreshTimestamp := lastRefreshTimestamp?
-          | let cancelled ← sleepWithCancellation timeUntilRefreshMs.toUInt32
+          | let cancelled ← sleepWithCancellation refreshIntervalMs.toUInt32
             if cancelled then
               return
             continue
 
         let currentTimestamp ← IO.monoMsNow
         let passedTimeMs := currentTimestamp - lastRefreshTimestamp
-        let remainingTimeMs := timeUntilRefreshMs - passedTimeMs
+        let remainingTimeMs := refreshIntervalMs - passedTimeMs
         if remainingTimeMs > 0 then
           let cancelled ← sleepWithCancellation remainingTimeMs.toUInt32
           if cancelled then
@@ -799,11 +820,10 @@ def runRefreshTasks : WorkerM (Array (Task Unit)) := do
           }
           (true, h)
         if ! canRefresh then
-          let cancelled ← sleepWithCancellation timeUntilRefreshMs.toUInt32
+          let cancelled ← sleepWithCancellation refreshIntervalMs.toUInt32
           if cancelled then
             return
           continue
-
         sendServerRequest ctx refreshMethod (none : Option Nat)
   return tasks
 

src/Lean/Server/FileWorker/InlayHints.lean

@@ -95,25 +95,30 @@ def applyEditToHint? (hintMod : Name) (ihi : Elab.InlayHintInfo) (range : String
   }
 
 structure InlayHintState where
-  oldInlayHints : Array Elab.InlayHintInfo
+  oldInlayHints      : Array Elab.InlayHintInfo
+  lastEditTimestamp? : Option Nat
   deriving TypeName, Inhabited
 
 def InlayHintState.init : InlayHintState := {
   oldInlayHints := #[]
+  lastEditTimestamp? := none
 }
 
 def handleInlayHints (_ : InlayHintParams) (s : InlayHintState) :
     RequestM (LspResponse (Array InlayHint) × InlayHintState) := do
   let ctx ← read
   let srcSearchPath := ctx.srcSearchPath
-  -- We delay sending inlay hints by 1000ms to avoid inlay hint flickering on the client.
+  -- We delay sending inlay hints by 3000ms to avoid inlay hint flickering on the client.
   -- VS Code already has a mechanism for this, but it is not sufficient.
-  -- Note that 1000ms of latency for this request are actually 2000ms of latency in VS Code after a
-  -- `textDocument/didChange` notification because VS Code (for some reason) emits two inlay hint
-  -- requests in succession after a change, immediately invalidating the result of the first.
-  -- Finally, for some stupid reason, VS Code doesn't remove the inlay hint when applying it,
-  -- so this additional latency causes the applied inlay hint to linger around for a bit.
-  let (snaps, _, isComplete) ← ctx.doc.cmdSnaps.getFinishedPrefixWithConsistentLatency 1000
+  let inlayHintEditDelayMs := 3000
+  let timestamp ← IO.monoMsNow
+  let editDelayMs :=
+    match s.lastEditTimestamp? with
+    | none => 0
+    | some lastEditTimestamp =>
+      let timeSinceLastEditMs := timestamp - lastEditTimestamp
+      inlayHintEditDelayMs - timeSinceLastEditMs
+  let (snaps, _, isComplete) ← ctx.doc.cmdSnaps.getFinishedPrefixWithConsistentLatency editDelayMs.toUInt32 (cancelTk? := ctx.cancelTk.truncatedTask)
   let finishedRange? : Option String.Range := do
     return ⟨⟨0⟩, ← List.max? <| snaps.map (fun s => s.endPos)⟩
   let oldInlayHints :=
@@ -133,7 +138,10 @@ def handleInlayHints (_ : InlayHintParams) (s : InlayHintState) :
         modify (·.push ih.toInlayHintInfo))
   let inlayHints := newInlayHints ++ oldInlayHints
   let lspInlayHints ← inlayHints.mapM (·.toLspInlayHint srcSearchPath ctx.doc.meta.text)
-  return ({ response := lspInlayHints, isComplete }, { s with oldInlayHints := inlayHints })
+  let r := { response := lspInlayHints, isComplete }
+  let s := { s with oldInlayHints := inlayHints }
+  RequestM.checkCanceled
+  return (r, s)
 
 def handleInlayHintsDidChange (p : DidChangeTextDocumentParams)
     : StateT InlayHintState RequestM Unit := do
@@ -164,12 +172,31 @@ def handleInlayHintsDidChange (p : DidChangeTextDocumentParams)
       ihi := ihi'
     if ! inlayHintInvalidated then
       updatedOldInlayHints := updatedOldInlayHints.push ihi
-  set <| { s with oldInlayHints := updatedOldInlayHints }
+  let isInlayHintInsertionEdit := p.contentChanges.all fun c => Id.run do
+    let .rangeChange changeRange newText := c
+      | return false
+    let changeRange := text.lspRangeToUtf8Range changeRange
+    let edit := ⟨changeRange, newText⟩
+    return s.oldInlayHints.any (·.textEdits.contains edit)
+  let timestamp ← IO.monoMsNow
+  let lastEditTimestamp? :=
+    if isInlayHintInsertionEdit then
+      -- For some stupid reason, VS Code doesn't remove the inlay hint when applying it, so we
+      -- try to figure out whether the edit was an insertion of an inlay hint and then respond
+      -- to the request without latency so that it inserted ASAP.
+      none
+    else
+      some timestamp
+  set <| { s with
+    oldInlayHints := updatedOldInlayHints
+    lastEditTimestamp?
+  }
 
 builtin_initialize
   registerPartialStatefulLspRequestHandler
     "textDocument/inlayHint"
     "workspace/inlayHint/refresh"
+    500
     InlayHintParams
     (Array InlayHint)
     InlayHintState