feat: highlight corresponding do (if any) when hovering over return

src/Lean/Data/Lsp/Capabilities.lean

@@ -30,6 +30,7 @@ instance ClientCapabilities.hasToJson : ToJson ClientCapabilities :=
 structure ServerCapabilities where
   textDocumentSync? : Option TextDocumentSyncOptions := none
   hoverProvider : Bool := false
+  documentHighlightProvider : Bool := false
   documentSymbolProvider : Bool := false
   definitionProvider : Bool := false
   declarationProvider : Bool := false

src/Lean/Data/Lsp/LanguageFeatures.lean

@@ -31,6 +31,27 @@ structure DefinitionParams extends TextDocumentPositionParams
 structure TypeDefinitionParams extends TextDocumentPositionParams
   deriving FromJson, ToJson
 
+structure DocumentHighlightParams extends TextDocumentPositionParams
+  deriving FromJson, ToJson
+
+inductive DocumentHighlightKind where
+  | text
+  | read
+  | write
+
+instance : ToJson DocumentHighlightKind where
+ toJson
+	| DocumentHighlightKind.text => 1
+	| DocumentHighlightKind.read => 2
+	| DocumentHighlightKind.write => 3
+
+structure DocumentHighlight where
+  range : Range
+  kind? : Option DocumentHighlightKind := none
+  deriving ToJson
+
+abbrev DocumentHighlightResult := Array DocumentHighlight
+
 structure DocumentSymbolParams where
   textDocument : TextDocumentIdentifier
   deriving FromJson, ToJson

src/Lean/Server/FileSource.lean

@@ -55,6 +55,9 @@ instance TypeDefinitionParams.hasFileSource : FileSource TypeDefinitionParams :=
 instance WaitForDiagnosticsParams.hasFileSource : FileSource WaitForDiagnosticsParams :=
   ⟨fun p => p.uri⟩
 
+instance DocumentHighlightParams.hasFileSource : FileSource DocumentHighlightParams :=
+  ⟨fun h => fileSource h.toTextDocumentPositionParams⟩
+
 instance DocumentSymbolParams.hasFileSource : FileSource DocumentSymbolParams :=
   ⟨fun p => fileSource p.textDocument⟩
 

src/Lean/Server/FileWorker.lean

@@ -453,6 +453,26 @@ section RequestHandling
     ⟨text.utf8PosToLspPos <| stx.getPos?.get!,
      text.utf8PosToLspPos <| stx.getTailPos?.get!⟩
 
+  partial def handleDocumentHighlight (p : DocumentHighlightParams) :
+    ServerM (Task (Except IO.Error (Except RequestError (Array DocumentHighlight)))) := do
+    let doc ← (← read).docRef.get
+    let text := doc.meta.text
+    let pos := text.lspPosToUtf8Pos p.position
+    let rec highlightReturn? (doRange? : Option Range) : Syntax → Option DocumentHighlight
+      | stx@`(doElem|return%$i $e) =>
+        if stx.getPos?.get! <= pos && pos < stx.getTailPos?.get! then
+          some { range := doRange?.getD (rangeOfSyntax text i), kind? := DocumentHighlightKind.text }
+        else
+          highlightReturn? doRange? e
+      | `(do%$i $elems) => highlightReturn? (rangeOfSyntax text i) elems
+      | stx => stx.getArgs.findSome? (highlightReturn? doRange?)
+
+    withWaitFindSnap doc (fun s => s.endPos > pos)
+      (notFoundX := pure #[]) fun snap => do
+        if let some hi := highlightReturn? none snap.stx then
+          return #[hi]
+        return #[]
+
   partial def handleDocumentSymbol (p : DocumentSymbolParams) :
     ServerM (Task (Except IO.Error (Except RequestError DocumentSymbolResult))) := do
     let st ← read
@@ -567,6 +587,7 @@ section MessageHandling
     | "textDocument/declaration"        => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := false)
     | "textDocument/definition"         => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := false)
     | "textDocument/typeDefinition"     => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := true)
+    | "textDocument/documentHighlight"  => handle DocumentHighlightParams DocumentHighlightResult handleDocumentHighlight
     | "textDocument/documentSymbol"     => handle DocumentSymbolParams DocumentSymbolResult handleDocumentSymbol
     | "$/lean/plainGoal"                => handle PlainGoalParams (Option PlainGoal) handlePlainGoal
     | _ => throwServerError s!"Got unsupported request: {method}"

src/Lean/Server/Watchdog.lean

@@ -386,6 +386,7 @@ section MessageHandling
     | "textDocument/declaration"        => handle DeclarationParams
     | "textDocument/definition"         => handle DefinitionParams
     | "textDocument/typeDefinition"     => handle TypeDefinitionParams
+    | "textDocument/documentHighlight"  => handle DocumentHighlightParams
     | "textDocument/documentSymbol"     => handle DocumentSymbolParams
     | "$/lean/plainGoal"                => handle PlainGoalParams
     | _                                 =>
@@ -497,6 +498,7 @@ def mkLeanServerCapabilities : ServerCapabilities := {
   declarationProvider := true
   definitionProvider := true
   typeDefinitionProvider := true
+  documentHighlightProvider := true
   documentSymbolProvider := true
 }