refactor: extend Lsp.ModuleRefs in Server.References

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -214,7 +214,7 @@ partial def handleDocumentHighlight (p : DocumentHighlightParams)
 
   let highlightRefs? (snaps : Array Snapshot) (pos : Lsp.Position) : Option (Array DocumentHighlight) := Id.run do
     let trees := snaps.map (·.infoTree)
-    let refs := findModuleRefs text trees
+    let refs : Lsp.ModuleRefs := findModuleRefs text trees
     let mut ranges := #[]
     for ident in ← refs.findAt p.position do
       if let some info ← refs.find? ident then

src/Lean/Server/References.lean

@@ -16,7 +16,7 @@ import Lean.Server.Snapshots
 /- Representing collected and deduplicated definitions and usages -/
 
 namespace Lean.Server
-open Lsp Lean.Elab
+open Lsp Lean.Elab Std
 
 structure Reference where
   ident : RefIdent
@@ -27,6 +27,43 @@ structure Reference where
   isBinder : Bool
   deriving Inhabited
 
+structure RefInfo where
+  definition : Option Reference
+  usages : Array Reference
+
+namespace RefInfo
+
+def empty : RefInfo := ⟨ none, #[] ⟩
+
+def addRef : RefInfo → Reference → RefInfo
+  | i@{ definition := none, .. }, ref@{ isBinder := true, .. } =>
+    { i with definition := ref }
+  | i@{ usages, .. }, ref@{ isBinder := false, .. } =>
+    { i with usages := usages.push ref }
+  | i, _ => i
+
+instance : Coe RefInfo Lsp.RefInfo where
+  coe self :=
+  {
+    definition := self.definition.map (·.range)
+    usages := self.usages.map (·.range)
+  }
+
+end RefInfo
+
+def ModuleRefs := HashMap RefIdent RefInfo
+
+namespace ModuleRefs
+
+def addRef (self : ModuleRefs) (ref : Reference) : ModuleRefs :=
+  let refInfo := self.findD ref.ident RefInfo.empty
+  self.insert ref.ident (refInfo.addRef ref)
+
+instance : Coe ModuleRefs Lsp.ModuleRefs where
+  coe self := HashMap.ofList <| List.map (fun (k, v) => (k, v)) <| self.toList
+
+end ModuleRefs
+
 end Lean.Server
 
 namespace Lean.Lsp.RefInfo
@@ -40,13 +77,6 @@ def merge (a : RefInfo) (b : RefInfo) : RefInfo :=
     usages := a.usages.append b.usages
   }
 
-def addRef : RefInfo → Reference → RefInfo
-  | i@{ definition := none, .. }, { range, isBinder := true, .. } =>
-    { i with definition := range }
-  | i@{ usages, .. }, { range, isBinder := false, .. } =>
-    { i with usages := usages.push range }
-  | i, _ => i
-
 def contains (self : RefInfo) (pos : Lsp.Position) : Bool := Id.run do
   if let some range := self.definition then
     if contains range pos then
@@ -64,10 +94,6 @@ end Lean.Lsp.RefInfo
 namespace Lean.Lsp.ModuleRefs
 open Server
 
-def addRef (self : ModuleRefs) (ref : Reference) : ModuleRefs :=
-  let refInfo := self.findD ref.ident RefInfo.empty
-  self.insert ref.ident (refInfo.addRef ref)
-
 def findAt (self : ModuleRefs) (pos : Lsp.Position) : Array RefIdent := Id.run do
   let mut result := #[]
   for (ident, info) in self.toList do
@@ -87,7 +113,7 @@ open Elab
 structure Ilean where
   version : Nat := 1
   module : Name
-  references : ModuleRefs
+  references : Lsp.ModuleRefs
   deriving FromJson, ToJson
 
 namespace Ilean
@@ -190,9 +216,9 @@ def findModuleRefs (text : FileMap) (trees : Array InfoTree) (localVars : Bool :
 
 structure References where
   /-- References loaded from ilean files -/
-  ileans : HashMap Name (System.FilePath × ModuleRefs)
+  ileans : HashMap Name (System.FilePath × Lsp.ModuleRefs)
   /-- References from workers, overriding the corresponding ilean files -/
-  workers : HashMap Name (Nat × ModuleRefs)
+  workers : HashMap Name (Nat × Lsp.ModuleRefs)
 
 namespace References
 
@@ -207,18 +233,18 @@ def removeIlean (self : References) (path : System.FilePath) : References :=
   namesToRemove.foldl (init := self) fun self name =>
     { self with ileans := self.ileans.erase name }
 
-def updateWorkerRefs (self : References) (name : Name) (version : Nat) (refs : ModuleRefs) : References := Id.run do
+def updateWorkerRefs (self : References) (name : Name) (version : Nat) (refs : Lsp.ModuleRefs) : References := Id.run do
   if let some (currVersion, _) := self.workers.find? name then
     if version > currVersion then
       return { self with workers := self.workers.insert name (version, refs) }
     if version == currVersion then
       let current := self.workers.findD name (version, HashMap.empty)
       let merged := refs.fold (init := current.snd) fun m ident info =>
-        m.findD ident RefInfo.empty |>.merge info |> m.insert ident
+        m.findD ident Lsp.RefInfo.empty |>.merge info |> m.insert ident
       return { self with workers := self.workers.insert name (version, merged) }
   return self
 
-def finalizeWorkerRefs (self : References) (name : Name) (version : Nat) (refs : ModuleRefs) : References := Id.run do
+def finalizeWorkerRefs (self : References) (name : Name) (version : Nat) (refs : Lsp.ModuleRefs) : References := Id.run do
   if let some (currVersion, _) := self.workers.find? name then
     if version < currVersion then
       return self
@@ -227,7 +253,7 @@ def finalizeWorkerRefs (self : References) (name : Name) (version : Nat) (refs :
 def removeWorkerRefs (self : References) (name : Name) : References :=
   { self with workers := self.workers.erase name }
 
-def allRefs (self : References) : HashMap Name ModuleRefs :=
+def allRefs (self : References) : HashMap Name Lsp.ModuleRefs :=
   let ileanRefs := self.ileans.toList.foldl (init := HashMap.empty) fun m (name, _, refs) => m.insert name refs
   self.workers.toList.foldl (init := ileanRefs) fun m (name, _, refs) => m.insert name refs