feat: inlay hint refinements

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