Add code lenses to definitions showing amount of references

src/Lean/Data/Lsp/Capabilities.lean

@@ -40,6 +40,7 @@ structure ServerCapabilities where
   hoverProvider : Bool := false
   documentHighlightProvider : Bool := false
   documentSymbolProvider : Bool := false
+  codeLensProvider? : Option CodeLensOptions := none
   definitionProvider : Bool := false
   declarationProvider : Bool := false
   typeDefinitionProvider : Bool := false

src/Lean/Data/Lsp/LanguageFeatures.lean

@@ -210,6 +210,27 @@ structure DocumentSymbolResult where
 instance : ToJson DocumentSymbolResult where
   toJson dsr := toJson dsr.syms
 
+structure CodeLensOptions where
+  resolveProvider? : Option Bool := false
+  deriving FromJson, ToJson
+
+structure CodeLensParams where
+  textDocument : TextDocumentIdentifier
+  deriving FromJson, ToJson
+
+inductive CodeLensInfo where
+  | ref (name : Name)
+  deriving FromJson, ToJson
+
+structure CodeLens where
+  range : Range
+  command? : Option Command
+  -- According to the LSP spec, the data field is optional and may contain any
+  -- sort of value. However, it is set by the server and preserved between
+  -- requests, so we can use a more specific type since we control its values.
+  data : CodeLensInfo
+  deriving FromJson, ToJson
+
 inductive SymbolTag where
   | deprecated
 

src/Lean/Server/References.lean

@@ -194,6 +194,16 @@ def findAt? (self : References) (module : Name) (pos : Lsp.Position) : Option Re
     return refs.findAt? pos
   none
 
+def countReferencesTo (self : References) (ident : RefIdent)
+    (includeDefinition : Bool := true) : Nat := Id.run do
+  let mut result := 0
+  for (module, refs) in self.allRefs.toList do
+    if let some info := refs.find? ident then
+      result := result + info.usages.size
+      if includeDefinition && info.definition.isSome then
+        result := result + 1
+  return result
+
 def referringTo (self : References) (ident : RefIdent) (srcSearchPath : SearchPath)
     (includeDefinition : Bool := true) : IO (Array Location) := do
   let mut result := #[]

src/Lean/Server/Watchdog.lean

@@ -213,6 +213,7 @@ section ServerM
     if let some path := fw.doc.meta.uri.toPath? then
       if let some module ← searchModuleNameOfFileName path s.srcSearchPath then
         s.references.modify fun refs => refs.addWorkerRefs module params.version params.references
+        s.hOut.writeLspNotification { method := "workspace/codeLens/refresh", param := Json.null }
 
   /-- Creates a Task which forwards a worker's messages into the output stream until an event
   which must be handled in the main watchdog thread (e.g. an I/O error) happens. -/
@@ -400,6 +401,47 @@ where
       targetIt := targetIt.next
     return false
 
+open Std in
+def handleCodeLens (p : CodeLensParams) : ServerM (Array CodeLens) := do
+  if let some path := p.textDocument.uri.toPath? then
+    let srcSearchPath := (← read).srcSearchPath
+    if let some module ← searchModuleNameOfFileName path srcSearchPath then
+      let references ← (← read).references.get
+      if let some refs := references.allRefs.find? module then
+        return lineDefs refs |>.map fun (name, range) =>
+          { range, command? := none, data := CodeLensInfo.ref name : CodeLens }
+  #[]
+where
+  lineDefs (refs : ModuleRefs) : Array (Name × Lsp.Range) := Id.run do
+    -- Construct hash map from line to definition. If there exists an entry in
+    -- the hash map, there's at least one definition for this line. Definitions
+    -- with local identifiers and multi-line definitions are ignored.
+    let lineMap := refs.fold (init := HashMap.empty) fun m ident info => Id.run do
+      if let (RefIdent.const name, some range) := (ident, info.definition) then
+        if range.start.line == range.end.line then
+          let line := range.start.line
+          if m.contains line then
+            -- Multiple definitions per line aren't allowed, but this line needs
+            -- to be marked as containing at least one definition.
+            return m.insert line none
+          else
+            return m.insert line $ some (name, range)
+      return m
+    return lineMap.fold (init := #[]) fun l _ info =>
+      if let some info := info then l.push info else l
+
+def handleCodeLensResolve (lens : CodeLens) : ServerM CodeLens := do
+  match lens.data with | CodeLensInfo.ref name => do
+    let references ← (← read).references.get
+    let count := references.countReferencesTo (RefIdent.const name) (includeDefinition := false)
+    let title := if count == 1 then s!"{count} reference" else s!"{count} references"
+    -- This needs to be implemented in the editor-specific lean plugins. The LSP
+    -- spec says: "The protocol currently doesn’t specify a set of well-known
+    -- commands." At least in VSCode, an empty command string makes the code
+    -- lenses non-clickable.
+    let command := ""
+    return { lens with command? := some { title, command } }
+
 end RequestHandling
 
 section NotificationHandling
@@ -517,6 +559,8 @@ section MessageHandling
     match method with
       | "textDocument/references" => handle ReferenceParams (Array Location) handleReference
       | "workspace/symbol" => handle WorkspaceSymbolParams (Array SymbolInformation) handleWorkspaceSymbol
+      | "textDocument/codeLens" => handle CodeLensParams (Array CodeLens) handleCodeLens
+      | "codeLens/resolve" => handle CodeLens CodeLens handleCodeLensResolve
       | _ => forwardRequestToWorker id method params
 
   def handleNotification (method : String) (params : Json) : ServerM Unit := do
@@ -655,6 +699,9 @@ def mkLeanServerCapabilities : ServerCapabilities := {
   workspaceSymbolProvider := true
   documentHighlightProvider := true
   documentSymbolProvider := true
+  codeLensProvider? := some {
+    resolveProvider? := true
+  }
   semanticTokensProvider? := some {
     legend := {
       tokenTypes     := SemanticTokenType.names