feat: add call hierarchy support

src/Lean/Data/HashMap.lean

@@ -227,3 +227,13 @@ def ofListWith (l : List (α × β)) (f : β → β → β) : HashMap α β :=
       match m.find? p.fst with
         | none   => m.insert p.fst p.snd
         | some v => m.insert p.fst $ f v p.snd)
+end Lean.HashMap
+
+def Array.groupByKey [BEq α] [Hashable α] (key : β → α) (xs : Array β)
+    : Lean.HashMap α (Array β) := Id.run do
+  let mut groups := ∅
+  for x in xs do
+    let group := groups.findD (key x) #[]
+    groups := groups.erase (key x) -- make `group` referentially unique
+    groups := groups.insert (key x) (group.push x)
+  return groups

src/Lean/Data/Lsp/Capabilities.lean

@@ -74,6 +74,7 @@ structure ServerCapabilities where
   declarationProvider       : Bool                           := false
   typeDefinitionProvider    : Bool                           := false
   referencesProvider        : Bool                           := false
+  callHierarchyProvider     : Bool                           := false
   renameProvider?           : Option RenameOptions           := none
   workspaceSymbolProvider   : Bool                           := false
   foldingRangeProvider      : Bool                           := false

src/Lean/Data/Lsp/Internal.lean

@@ -19,7 +19,7 @@ reduce the size of the resulting JSON. -/
 
 inductive RefIdent where
   | const : Name → RefIdent
-  | fvar : FVarId → RefIdent
+  | fvar  : FVarId → RefIdent
   deriving BEq, Hashable, Inhabited
 
 namespace RefIdent
@@ -43,33 +43,79 @@ def fromString (s : String) : Except String RefIdent := do
     | "f:" => return RefIdent.fvar <| FVarId.mk name
     | _ => throw "string must start with 'c:' or 'f:'"
 
+instance : FromJson RefIdent where
+  fromJson?
+    | (s : String) => fromString s
+    | j => Except.error s!"expected a String, got {j}"
+
+instance : ToJson RefIdent where
+  toJson ident := toString ident
+
 end RefIdent
 
+structure RefInfo.ParentDecl where
+  name           : Name
+  range          : Lsp.Range
+  selectionRange : Lsp.Range
+  deriving ToJson
+
+structure RefInfo.Location where
+  range       : Lsp.Range
+  parentDecl? : Option RefInfo.ParentDecl
+
 structure RefInfo where
-  definition : Option Lsp.Range
-  usages : Array Lsp.Range
+  definition? : Option RefInfo.Location
+  usages      : Array RefInfo.Location
 
 instance : ToJson RefInfo where
   toJson i :=
     let rangeToList (r : Lsp.Range) : List Nat :=
       [r.start.line, r.start.character, r.end.line, r.end.character]
+    let parentDeclToList (d : RefInfo.ParentDecl) : List Json :=
+      let name := d.name.toString |> toJson
+      let range := rangeToList d.range |>.map toJson
+      let selectionRange := rangeToList d.selectionRange |>.map toJson
+      [name] ++ range ++ selectionRange
+    let locationToList (l : RefInfo.Location) : List Json :=
+      let range := rangeToList l.range |>.map toJson
+      let parentDecl := l.parentDecl?.map parentDeclToList |>.getD []
+      range ++ parentDecl
     Json.mkObj [
-      ("definition", toJson $ i.definition.map rangeToList),
-      ("usages", toJson $ i.usages.map rangeToList)
+      ("definition", toJson $ i.definition?.map locationToList),
+      ("usages", toJson $ i.usages.map locationToList)
     ]
 
 instance : FromJson RefInfo where
   fromJson? j := do
-    let listToRange (l : List Nat) : Except String Lsp.Range := match l with
+    let toRange : List Nat → Except String Lsp.Range
       | [sLine, sChar, eLine, eChar] => pure ⟨⟨sLine, sChar⟩, ⟨eLine, eChar⟩⟩
-      | _ => throw s!"Expected list of length 4, not {l.length}"
-    let definition ← j.getObjValAs? (Option $ List Nat) "definition"
-    let definition ← match definition with
+      | l => throw s!"Expected list of length 4, not {l.length}"
+    let toParentDecl (a : Array Json) : Except String RefInfo.ParentDecl := do
+      let name := String.toName <| ← fromJson? a[0]!
+      let range ← a[1:5].toArray.toList |>.mapM fromJson?
+      let range ← toRange range
+      let selectionRange ← a[5:].toArray.toList |>.mapM fromJson?
+      let selectionRange ← toRange selectionRange
+      return ⟨name, range, selectionRange⟩
+    let toLocation (l : List Json) : Except String RefInfo.Location := do
+      let l := l.toArray
+      if l.size != 4 && l.size != 13 then
+        .error "Expected list of length 4 or 13, not {l.size}"
+      let range ← l[:4].toArray.toList |>.mapM fromJson?
+      let range ← toRange range
+      if l.size == 13 then
+        let parentDecl ← toParentDecl l[4:].toArray
+        return ⟨range, parentDecl⟩
+      else
+        return ⟨range, none⟩
+
+    let definition? ← j.getObjValAs? (Option $ List Json) "definition"
+    let definition? ← match definition? with
       | none => pure none
-      | some list => some <$> listToRange list
-    let usages ← j.getObjValAs? (Array $ List Nat) "usages"
-    let usages ← usages.mapM listToRange
-    pure { definition, usages }
+      | some list => some <$> toLocation list
+    let usages ← j.getObjValAs? (Array $ List Json) "usages"
+    let usages ← usages.mapM toLocation
+    pure { definition?, usages }
 
 /-- References from a single module/file -/
 def ModuleRefs := HashMap RefIdent RefInfo
@@ -88,7 +134,7 @@ instance : FromJson ModuleRefs where
 Contains the file's definitions and references. -/
 structure LeanIleanInfoParams where
   /-- Version of the file these references are from. -/
-  version : Nat
+  version    : Nat
   references : ModuleRefs
   deriving FromJson, ToJson
 

src/Lean/Data/Lsp/LanguageFeatures.lean

@@ -153,35 +153,66 @@ inductive SymbolKind where
   | event
   | operator
   | typeParameter
+  deriving BEq, Hashable, Inhabited
+
+instance : FromJson SymbolKind where
+  fromJson?
+    | 1  => .ok .file
+    | 2  => .ok .module
+    | 3  => .ok .namespace
+    | 4  => .ok .package
+    | 5  => .ok .class
+    | 6  => .ok .method
+    | 7  => .ok .property
+    | 8  => .ok .field
+    | 9  => .ok .constructor
+    | 10 => .ok .enum
+    | 11 => .ok .interface
+    | 12 => .ok .function
+    | 13 => .ok .variable
+    | 14 => .ok .constant
+    | 15 => .ok .string
+    | 16 => .ok .number
+    | 17 => .ok .boolean
+    | 18 => .ok .array
+    | 19 => .ok .object
+    | 20 => .ok .key
+    | 21 => .ok .null
+    | 22 => .ok .enumMember
+    | 23 => .ok .struct
+    | 24 => .ok .event
+    | 25 => .ok .operator
+    | 26 => .ok .typeParameter
+    | j  => .error s!"invalid symbol kind {j}"
 
 instance : ToJson SymbolKind where
- toJson
-   | SymbolKind.file => 1
-   | SymbolKind.module => 2
-   | SymbolKind.namespace => 3
-   | SymbolKind.package => 4
-   | SymbolKind.class => 5
-   | SymbolKind.method => 6
-   | SymbolKind.property => 7
-   | SymbolKind.field => 8
-   | SymbolKind.constructor => 9
-   | SymbolKind.enum => 10
-   | SymbolKind.interface => 11
-   | SymbolKind.function => 12
-   | SymbolKind.variable => 13
-   | SymbolKind.constant => 14
-   | SymbolKind.string => 15
-   | SymbolKind.number => 16
-   | SymbolKind.boolean => 17
-   | SymbolKind.array => 18
-   | SymbolKind.object => 19
-   | SymbolKind.key => 20
-   | SymbolKind.null => 21
-   | SymbolKind.enumMember => 22
-   | SymbolKind.struct => 23
-   | SymbolKind.event => 24
-   | SymbolKind.operator => 25
-   | SymbolKind.typeParameter => 26
+  toJson
+    | .file          => 1
+    | .module        => 2
+    | .namespace     => 3
+    | .package       => 4
+    | .class         => 5
+    | .method        => 6
+    | .property      => 7
+    | .field         => 8
+    | .constructor   => 9
+    | .enum          => 10
+    | .interface     => 11
+    | .function      => 12
+    | .variable      => 13
+    | .constant      => 14
+    | .string        => 15
+    | .number        => 16
+    | .boolean       => 17
+    | .array         => 18
+    | .object        => 19
+    | .key           => 20
+    | .null          => 21
+    | .enumMember    => 22
+    | .struct        => 23
+    | .event         => 24
+    | .operator      => 25
+    | .typeParameter => 26
 
 structure DocumentSymbolAux (Self : Type) where
   name           : String
@@ -191,7 +222,7 @@ structure DocumentSymbolAux (Self : Type) where
   range          : Range
   selectionRange : Range
   children?      : Option (Array Self) := none
-  deriving ToJson
+  deriving FromJson, ToJson
 
 inductive DocumentSymbol where
   | mk (sym : DocumentSymbolAux DocumentSymbol)
@@ -212,18 +243,56 @@ instance : ToJson DocumentSymbolResult where
 
 inductive SymbolTag where
   | deprecated
+  deriving BEq, Hashable, Inhabited
+
+instance : FromJson SymbolTag where
+  fromJson?
+    | 1 => .ok .deprecated
+    | j => .error s!"unknown symbol tag {j}"
 
 instance : ToJson SymbolTag where
- toJson
-   | SymbolTag.deprecated => 1
+  toJson
+    | .deprecated => 1
 
 structure SymbolInformation where
   name           : String
   kind           : SymbolKind
   tags           : Array SymbolTag := #[]
   location       : Location
   containerName? : Option String := none
-  deriving ToJson
+  deriving FromJson, ToJson
+
+structure CallHierarchyPrepareParams extends TextDocumentPositionParams
+  deriving FromJson, ToJson
+
+structure CallHierarchyItem where
+  name           : String
+  kind           : SymbolKind
+  tags?          : Option (Array SymbolTag) := none
+  detail?        : Option String := none
+  uri            : DocumentUri
+  range          : Range
+  selectionRange : Range
+  -- data? : Option unknown
+  deriving FromJson, ToJson, BEq, Hashable, Inhabited
+
+structure CallHierarchyIncomingCallsParams where
+  item : CallHierarchyItem
+  deriving FromJson, ToJson
+
+structure CallHierarchyIncomingCall where
+  «from»     : CallHierarchyItem
+  fromRanges : Array Range
+  deriving FromJson, ToJson, Inhabited
+
+structure CallHierarchyOutgoingCallsParams where
+  item : CallHierarchyItem
+  deriving FromJson, ToJson
+
+structure CallHierarchyOutgoingCall where
+  to         : CallHierarchyItem
+  fromRanges : Array Range
+  deriving FromJson, ToJson, Inhabited
 
 inductive SemanticTokenType where
   -- Used by Lean

src/Lean/Data/Lsp/Utf16.lean

@@ -8,6 +8,7 @@ import Init.Data.String
 import Init.Data.Array
 import Lean.Data.Lsp.Basic
 import Lean.Data.Position
+import Lean.DeclarationRange
 
 /-! LSP uses UTF-16 for indexing, so we need to provide some primitives
 to interact with Lean strings using UTF-16 indices. -/
@@ -86,3 +87,8 @@ def utf8PosToLspPos (text : FileMap) (pos : String.Pos) : Lsp.Position :=
 
 end FileMap
 end Lean
+
+def Lean.DeclarationRange.toLspRange (r : Lean.DeclarationRange) : Lsp.Range := {
+  start := ⟨r.pos.line - 1, r.charUtf16⟩
+  «end» := ⟨r.endPos.line - 1, r.endCharUtf16⟩
+}

src/Lean/Elab/Frontend.lean

@@ -118,7 +118,7 @@ def runFrontend
   if let some ileanFileName := ileanFileName? then
     let trees := s.commandState.infoState.trees.toArray
     let references := Lean.Server.findModuleRefs inputCtx.fileMap trees (localVars := false)
-    let ilean := { module := mainModuleName, references : Lean.Server.Ilean }
+    let ilean := { module := mainModuleName, references := ← references.toLspModuleRefs : Lean.Server.Ilean }
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
   pure (s.commandState.env, !s.commandState.messages.hasErrors)

src/Lean/Server/FileWorker.lean

@@ -76,7 +76,7 @@ section Elab
   private def publishIleanInfo (method : String) (m : DocumentMeta) (hOut : FS.Stream)
       (snaps : Array Snapshot) : IO Unit := do
     let trees := snaps.map fun snap => snap.infoTree
-    let references := findModuleRefs m.text trees (localVars := true)
+    let references ← findModuleRefs m.text trees (localVars := true) |>.toLspModuleRefs
     let param := { version := m.version, references : LeanIleanInfoParams }
     hOut.writeLspNotification { method, param }
 

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -290,23 +290,23 @@ partial def handleDocumentHighlight (p : DocumentHighlightParams)
     | `(do%$i $elems) => highlightReturn? (i.getRange?.get!.toLspRange text) elems
     | stx => stx.getArgs.findSome? (highlightReturn? doRange?)
 
-  let highlightRefs? (snaps : Array Snapshot) : Option (Array DocumentHighlight) := Id.run do
+  let highlightRefs? (snaps : Array Snapshot) : IO (Option (Array DocumentHighlight)) := do
     let trees := snaps.map (·.infoTree)
-    let refs : Lsp.ModuleRefs := findModuleRefs text trees
+    let refs : Lsp.ModuleRefs ← findModuleRefs text trees |>.toLspModuleRefs
     let mut ranges := #[]
-    for ident in ← refs.findAt p.position do
-      if let some info ← refs.find? ident then
-        if let some definition := info.definition then
-          ranges := ranges.push definition
-        ranges := ranges.append info.usages
+    for ident in refs.findAt p.position do
+      if let some info := refs.find? ident then
+        if let some ⟨definitionRange, _⟩ := info.definition? then
+          ranges := ranges.push definitionRange
+        ranges := ranges.append <| info.usages.map (·.range)
     if ranges.isEmpty then
       return none
-    some <| ranges.map ({ range := ·, kind? := DocumentHighlightKind.text })
+    return some <| ranges.map ({ range := ·, kind? := DocumentHighlightKind.text })
 
   withWaitFindSnap doc (fun s => s.endPos > pos)
     (notFoundX := pure #[]) fun snap => do
       let (snaps, _) ← doc.cmdSnaps.getFinishedPrefix
-      if let some his := highlightRefs? snaps.toArray then
+      if let some his ← highlightRefs? snaps.toArray then
         return his
       if let some hi := highlightReturn? none snap.stx then
         return #[hi]

src/Lean/Server/GoTo.lean

@@ -34,14 +34,11 @@ def locationLinksFromDecl (srcSearchPath : SearchPath) (uri : DocumentUri) (n :
 
   let ranges? ← findDeclarationRanges? n
   if let (some ranges, some modUri) := (ranges?, modUri?) then
-    let declRangeToLspRange (r : DeclarationRange) : Lsp.Range :=
-      { start := ⟨r.pos.line - 1, r.charUtf16⟩
-        «end» := ⟨r.endPos.line - 1, r.endCharUtf16⟩ }
     let ll : LocationLink := {
       originSelectionRange? := originRange?
       targetUri := modUri
-      targetRange := declRangeToLspRange ranges.range
-      targetSelectionRange := declRangeToLspRange ranges.selectionRange
+      targetRange := ranges.range.toLspRange
+      targetSelectionRange := ranges.selectionRange.toLspRange
     }
     return #[ll]
   return #[]