feat: load ilean files from olean search path

src/Lean/Elab/Frontend.lean

@@ -119,8 +119,8 @@ def runFrontend
 
   if let some ileanFileName := ileanFileName then
     let trees := s.commandState.infoState.trees.toList
-    let references := Lean.Server.References.findFileRefs inputCtx.fileMap trees (localVars := false)
-    let ilean := { module := mainModuleName, references : Lean.Server.References.Ilean }
+    let references := Lean.Server.findFileRefs inputCtx.fileMap trees (localVars := false)
+    let ilean := { module := mainModuleName, references : Lean.Server.Ilean }
     IO.FS.writeFile ileanFileName $ toString $ toJson ilean
 
   pure (s.commandState.env, !s.commandState.messages.hasErrors)

src/Lean/Server/References.lean

@@ -4,7 +4,7 @@ import Lean.Data.Lsp
 import Lean.Server.InfoUtils
 import Lean.Server.Snapshots
 
-namespace Lean.Server.References
+namespace Lean.Server
 open Std
 open Lsp
 open Elab
@@ -104,8 +104,6 @@ def addRef (self : FileRefMap) (ref : Reference) : FileRefMap :=
 
 end FileRefMap
 
-def RefMap := HashMap String FileRefMap
-
 /-- Content of individual `.ilean` files -/
 structure Ilean where
   version : Nat := 1
@@ -176,4 +174,36 @@ def findFileRefs (text : FileMap) (trees : List InfoTree) (localVars : Bool := t
       | _ => true
   refs.foldl (init := HashMap.empty) fun m ref => m.addRef ref
 
-end Lean.Server.References
+/- Collecting and maintaining reference info from different sources -/
+
+structure References where
+  /-- References loaded from ilean files -/
+  ileans : HashMap Name (System.FilePath × FileRefMap)
+  /-- References from workers, overriding the corresponding ilean files -/
+  workers : HashMap Name (Nat × FileRefMap)
+
+namespace References
+
+def empty : References := { ileans := HashMap.empty, workers := HashMap.empty }
+
+def addIlean (self : References) (path : System.FilePath) (ilean : Ilean) : References :=
+  { self with ileans := self.ileans.insert ilean.module (path, ilean.references) }
+
+def removeIlean (self : References) (path : System.FilePath) : References :=
+  let namesToRemove := self.ileans.toList.filter (fun (_, p, _) => p == path)
+    |>.map (fun (n, _, _) => n)
+  namesToRemove.foldl (init := self) fun self name =>
+    { self with ileans := self.ileans.erase name }
+
+def addWorkerRefs (self : References) (name : Name) (version : Nat) (refs : FileRefMap) : References := Id.run do
+  if let some (currVersion, _) := self.workers.find? name then
+    if version <= currVersion then
+      return self
+  return { self with workers := self.workers.insert name (version, refs) }
+
+def removeWorkerRefs (self : References) (name : Name) : References :=
+  { self with workers := self.workers.erase name }
+
+end References
+
+end Lean.Server

src/Lean/Server/Watchdog.lean

@@ -13,6 +13,7 @@ import Lean.Elab.Import
 import Lean.Data.Lsp
 import Lean.Server.Utils
 import Lean.Server.Requests
+import Lean.Server.References
 
 /-!
 For general server architecture, see `README.md`. This module implements the watchdog process.
@@ -178,6 +179,7 @@ section ServerM
     editDelay      : Nat
     workerPath     : System.FilePath
     srcSearchPath  : System.SearchPath
+    references     : References
 
   abbrev ServerM := ReaderT ServerContext IO
 
@@ -595,9 +597,21 @@ def findSrcSearchPath : IO System.SearchPath := do
     srcSearchPath := System.SearchPath.parse p ++ srcSearchPath
   srcSearchPath
 
+def loadReferences : IO References := do
+  let oleanSearchPath ← Lean.searchPathRef.get
+  let mut refs := References.empty
+  for path in ← oleanSearchPath.findAllWithExt "ilean" do
+    let content ← FS.readFile path
+    let ilean ← match Json.parse content >>= fromJson? with
+      | Except.ok ilean => pure ilean
+      | Except.error msg => throwServerError s!"Failed to load ilean at {path}: {msg}"
+    refs := refs.addIlean path ilean
+  refs
+
 def initAndRunWatchdog (args : List String) (i o e : FS.Stream) : IO Unit := do
   let workerPath ← findWorkerPath
   let srcSearchPath ← findSrcSearchPath
+  let references ← loadReferences
   let fileWorkersRef ← IO.mkRef (RBMap.empty : FileWorkerMap)
   let i ← maybeTee "wdIn.txt" false i
   let o ← maybeTee "wdOut.txt" true o
@@ -637,6 +651,7 @@ def initAndRunWatchdog (args : List String) (i o e : FS.Stream) : IO Unit := do
     editDelay      := initRequest.param.initializationOptions? |>.bind InitializationOptions.editDelay? |>.getD 200
     workerPath
     srcSearchPath
+    references
     : ServerContext
   }
 

src/Lean/Util/Path.lean

@@ -37,6 +37,14 @@ def findWithExt (sp : SearchPath) (ext : String) (mod : Name) : IO (Option FileP
     (p / pkg).isDir <||> ((p / pkg).withExtension ext).pathExists
   return root?.map (modToFilePath · mod ext)
 
+def findAllWithExt (sp : SearchPath) (ext : String) : IO (Array FilePath) := do
+  let mut paths := #[]
+  for p in sp do
+    for e in ← p.readDir do
+      if e.fileName.endsWith s!".{ext}" then
+        paths := paths.push e.path
+  paths
+
 end SearchPath
 
 builtin_initialize searchPathRef : IO.Ref SearchPath ← IO.mkRef {}